Skip to content

Engine API

odibi.engine.base

Base engine interface.

Engine

Bases: ABC

Abstract base class for execution engines.

Source code in odibi/engine/base.py
class Engine(ABC):
    """Abstract base class for execution engines."""

    # Custom format registry
    _custom_readers: Dict[str, Any] = {}
    _custom_writers: Dict[str, Any] = {}

    @classmethod
    def register_format(cls, fmt: str, reader: Optional[Any] = None, writer: Optional[Any] = None):
        """Register custom format reader/writer.

        Args:
            fmt: Format name (e.g. 'netcdf')
            reader: Function(path, **options) -> DataFrame
            writer: Function(df, path, **options) -> None
        """
        if reader:
            cls._custom_readers[fmt] = reader
        if writer:
            cls._custom_writers[fmt] = writer

    @abstractmethod
    def read(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Read data from source.

        Args:
            connection: Connection object
            format: Data format (csv, parquet, delta, etc.)
            table: Table name (for SQL/Delta)
            path: File path (for file-based sources)
            options: Format-specific options

        Returns:
            DataFrame (engine-specific type)
        """
        pass

    def materialize(self, df: Any) -> Any:
        """Materialize lazy dataset into memory (DataFrame).

        Args:
            df: DataFrame or LazyDataset

        Returns:
            Materialized DataFrame
        """
        return df

    @abstractmethod
    def write(
        self,
        df: Any,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        mode: str = "overwrite",
        options: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Any] = None,
    ) -> None:
        """Write data to destination.

        Args:
            df: DataFrame to write
            connection: Connection object
            format: Output format
            table: Table name (for SQL/Delta)
            path: File path (for file-based outputs)
            mode: Write mode (overwrite/append)
            options: Format-specific options
        """
        pass

    @abstractmethod
    def execute_sql(self, sql: str, context: Context) -> Any:
        """Execute SQL query.

        Args:
            sql: SQL query string
            context: Execution context with registered DataFrames

        Returns:
            Result DataFrame
        """
        pass

    @abstractmethod
    def execute_operation(self, operation: str, params: Dict[str, Any], df: Any) -> Any:
        """Execute built-in operation (pivot, etc.).

        Args:
            operation: Operation name
            params: Operation parameters
            df: Input DataFrame

        Returns:
            Result DataFrame
        """
        pass

    @abstractmethod
    def get_schema(self, df: Any) -> Any:
        """Get DataFrame schema.

        Args:
            df: DataFrame

        Returns:
            Dict[str, str] mapping column names to types, or List[str] of names (deprecated)
        """
        pass

    @abstractmethod
    def get_shape(self, df: Any) -> tuple:
        """Get DataFrame shape.

        Args:
            df: DataFrame

        Returns:
            (rows, columns)
        """
        pass

    @abstractmethod
    def count_rows(self, df: Any) -> int:
        """Count rows in DataFrame.

        Args:
            df: DataFrame

        Returns:
            Row count
        """
        pass

    @abstractmethod
    def count_nulls(self, df: Any, columns: List[str]) -> Dict[str, int]:
        """Count nulls in specified columns.

        Args:
            df: DataFrame
            columns: Columns to check

        Returns:
            Dictionary of column -> null count
        """
        pass

    @abstractmethod
    def validate_schema(self, df: Any, schema_rules: Dict[str, Any]) -> List[str]:
        """Validate DataFrame schema.

        Args:
            df: DataFrame
            schema_rules: Validation rules

        Returns:
            List of validation failures (empty if valid)
        """
        pass

    @abstractmethod
    def validate_data(self, df: Any, validation_config: Any) -> List[str]:
        """Validate data against rules.

        Args:
            df: DataFrame to validate
            validation_config: ValidationConfig object

        Returns:
            List of validation failure messages (empty if valid)
        """
        pass

    @abstractmethod
    def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
        """Get sample rows as list of dictionaries.

        Args:
            df: DataFrame
            n: Number of rows to return

        Returns:
            List of row dictionaries
        """
        pass

    def get_source_files(self, df: Any) -> List[str]:
        """Get list of source files that generated this DataFrame.

        Args:
            df: DataFrame

        Returns:
            List of file paths (or empty list if not applicable/supported)
        """
        return []

    def profile_nulls(self, df: Any) -> Dict[str, float]:
        """Calculate null percentage for each column.

        Args:
            df: DataFrame

        Returns:
            Dictionary of {column_name: null_percentage} (0.0 to 1.0)
        """
        return {}

    @abstractmethod
    def table_exists(
        self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
    ) -> bool:
        """Check if table or location exists.

        Args:
            connection: Connection object
            table: Table name (for catalog tables)
            path: File path (for path-based tables)

        Returns:
            True if table/location exists, False otherwise
        """
        pass

    @abstractmethod
    def harmonize_schema(self, df: Any, target_schema: Dict[str, str], policy: Any) -> Any:
        """Harmonize DataFrame schema with target schema according to policy.

        Args:
            df: Input DataFrame
            target_schema: Target schema (column name -> type)
            policy: SchemaPolicyConfig object

        Returns:
            Harmonized DataFrame
        """
        pass

    @abstractmethod
    def anonymize(
        self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
    ) -> Any:
        """Anonymize specified columns.

        Args:
            df: DataFrame to anonymize
            columns: List of columns to anonymize
            method: Method ('hash', 'mask', 'redact')
            salt: Optional salt for hashing

        Returns:
            Anonymized DataFrame
        """
        pass

    def get_table_schema(
        self,
        connection: Any,
        table: Optional[str] = None,
        path: Optional[str] = None,
        format: Optional[str] = None,
    ) -> Optional[Dict[str, str]]:
        """Get schema of an existing table/file.

        Args:
            connection: Connection object
            table: Table name
            path: File path
            format: Data format (optional, helps with file-based sources)

        Returns:
            Schema dict or None if table doesn't exist or schema fetch fails.
        """
        return None

    def maintain_table(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        config: Optional[Any] = None,
    ) -> None:
        """Run table maintenance operations (optimize, vacuum).

        Args:
            connection: Connection object
            format: Table format
            table: Table name
            path: Table path
            config: AutoOptimizeConfig object
        """
        pass

    def add_write_metadata(
        self,
        df: Any,
        metadata_config: Any,
        source_connection: Optional[str] = None,
        source_table: Optional[str] = None,
        source_path: Optional[str] = None,
        is_file_source: bool = False,
    ) -> Any:
        """Add metadata columns to DataFrame before writing (Bronze layer lineage).

        Args:
            df: DataFrame
            metadata_config: WriteMetadataConfig or True (for all defaults)
            source_connection: Name of the source connection
            source_table: Name of the source table (SQL sources)
            source_path: Path of the source file (file sources)
            is_file_source: True if source is a file-based read

        Returns:
            DataFrame with metadata columns added (or unchanged if metadata_config is None/False)
        """
        return df  # Default: no-op

add_write_metadata(df, metadata_config, source_connection=None, source_table=None, source_path=None, is_file_source=False)

Add metadata columns to DataFrame before writing (Bronze layer lineage).

Parameters:

Name Type Description Default
df Any

DataFrame

required
metadata_config Any

WriteMetadataConfig or True (for all defaults)

required
source_connection Optional[str]

Name of the source connection

None
source_table Optional[str]

Name of the source table (SQL sources)

None
source_path Optional[str]

Path of the source file (file sources)

None
is_file_source bool

True if source is a file-based read

False

Returns:

Type Description
Any

DataFrame with metadata columns added (or unchanged if metadata_config is None/False)

Source code in odibi/engine/base.py
def add_write_metadata(
    self,
    df: Any,
    metadata_config: Any,
    source_connection: Optional[str] = None,
    source_table: Optional[str] = None,
    source_path: Optional[str] = None,
    is_file_source: bool = False,
) -> Any:
    """Add metadata columns to DataFrame before writing (Bronze layer lineage).

    Args:
        df: DataFrame
        metadata_config: WriteMetadataConfig or True (for all defaults)
        source_connection: Name of the source connection
        source_table: Name of the source table (SQL sources)
        source_path: Path of the source file (file sources)
        is_file_source: True if source is a file-based read

    Returns:
        DataFrame with metadata columns added (or unchanged if metadata_config is None/False)
    """
    return df  # Default: no-op

anonymize(df, columns, method, salt=None) abstractmethod

Anonymize specified columns.

Parameters:

Name Type Description Default
df Any

DataFrame to anonymize

required
columns List[str]

List of columns to anonymize

required
method str

Method ('hash', 'mask', 'redact')

required
salt Optional[str]

Optional salt for hashing

None

Returns:

Type Description
Any

Anonymized DataFrame

Source code in odibi/engine/base.py
@abstractmethod
def anonymize(
    self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
) -> Any:
    """Anonymize specified columns.

    Args:
        df: DataFrame to anonymize
        columns: List of columns to anonymize
        method: Method ('hash', 'mask', 'redact')
        salt: Optional salt for hashing

    Returns:
        Anonymized DataFrame
    """
    pass

count_nulls(df, columns) abstractmethod

Count nulls in specified columns.

Parameters:

Name Type Description Default
df Any

DataFrame

required
columns List[str]

Columns to check

required

Returns:

Type Description
Dict[str, int]

Dictionary of column -> null count

Source code in odibi/engine/base.py
@abstractmethod
def count_nulls(self, df: Any, columns: List[str]) -> Dict[str, int]:
    """Count nulls in specified columns.

    Args:
        df: DataFrame
        columns: Columns to check

    Returns:
        Dictionary of column -> null count
    """
    pass

count_rows(df) abstractmethod

Count rows in DataFrame.

Parameters:

Name Type Description Default
df Any

DataFrame

required

Returns:

Type Description
int

Row count

Source code in odibi/engine/base.py
@abstractmethod
def count_rows(self, df: Any) -> int:
    """Count rows in DataFrame.

    Args:
        df: DataFrame

    Returns:
        Row count
    """
    pass

execute_operation(operation, params, df) abstractmethod

Execute built-in operation (pivot, etc.).

Parameters:

Name Type Description Default
operation str

Operation name

required
params Dict[str, Any]

Operation parameters

required
df Any

Input DataFrame

required

Returns:

Type Description
Any

Result DataFrame

Source code in odibi/engine/base.py
@abstractmethod
def execute_operation(self, operation: str, params: Dict[str, Any], df: Any) -> Any:
    """Execute built-in operation (pivot, etc.).

    Args:
        operation: Operation name
        params: Operation parameters
        df: Input DataFrame

    Returns:
        Result DataFrame
    """
    pass

execute_sql(sql, context) abstractmethod

Execute SQL query.

Parameters:

Name Type Description Default
sql str

SQL query string

required
context Context

Execution context with registered DataFrames

required

Returns:

Type Description
Any

Result DataFrame

Source code in odibi/engine/base.py
@abstractmethod
def execute_sql(self, sql: str, context: Context) -> Any:
    """Execute SQL query.

    Args:
        sql: SQL query string
        context: Execution context with registered DataFrames

    Returns:
        Result DataFrame
    """
    pass

get_sample(df, n=10) abstractmethod

Get sample rows as list of dictionaries.

Parameters:

Name Type Description Default
df Any

DataFrame

required
n int

Number of rows to return

10

Returns:

Type Description
List[Dict[str, Any]]

List of row dictionaries

Source code in odibi/engine/base.py
@abstractmethod
def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
    """Get sample rows as list of dictionaries.

    Args:
        df: DataFrame
        n: Number of rows to return

    Returns:
        List of row dictionaries
    """
    pass

get_schema(df) abstractmethod

Get DataFrame schema.

Parameters:

Name Type Description Default
df Any

DataFrame

required

Returns:

Type Description
Any

Dict[str, str] mapping column names to types, or List[str] of names (deprecated)

Source code in odibi/engine/base.py
@abstractmethod
def get_schema(self, df: Any) -> Any:
    """Get DataFrame schema.

    Args:
        df: DataFrame

    Returns:
        Dict[str, str] mapping column names to types, or List[str] of names (deprecated)
    """
    pass

get_shape(df) abstractmethod

Get DataFrame shape.

Parameters:

Name Type Description Default
df Any

DataFrame

required

Returns:

Type Description
tuple

(rows, columns)

Source code in odibi/engine/base.py
@abstractmethod
def get_shape(self, df: Any) -> tuple:
    """Get DataFrame shape.

    Args:
        df: DataFrame

    Returns:
        (rows, columns)
    """
    pass

get_source_files(df)

Get list of source files that generated this DataFrame.

Parameters:

Name Type Description Default
df Any

DataFrame

required

Returns:

Type Description
List[str]

List of file paths (or empty list if not applicable/supported)

Source code in odibi/engine/base.py
def get_source_files(self, df: Any) -> List[str]:
    """Get list of source files that generated this DataFrame.

    Args:
        df: DataFrame

    Returns:
        List of file paths (or empty list if not applicable/supported)
    """
    return []

get_table_schema(connection, table=None, path=None, format=None)

Get schema of an existing table/file.

Parameters:

Name Type Description Default
connection Any

Connection object

required
table Optional[str]

Table name

None
path Optional[str]

File path

None
format Optional[str]

Data format (optional, helps with file-based sources)

None

Returns:

Type Description
Optional[Dict[str, str]]

Schema dict or None if table doesn't exist or schema fetch fails.

Source code in odibi/engine/base.py
def get_table_schema(
    self,
    connection: Any,
    table: Optional[str] = None,
    path: Optional[str] = None,
    format: Optional[str] = None,
) -> Optional[Dict[str, str]]:
    """Get schema of an existing table/file.

    Args:
        connection: Connection object
        table: Table name
        path: File path
        format: Data format (optional, helps with file-based sources)

    Returns:
        Schema dict or None if table doesn't exist or schema fetch fails.
    """
    return None

harmonize_schema(df, target_schema, policy) abstractmethod

Harmonize DataFrame schema with target schema according to policy.

Parameters:

Name Type Description Default
df Any

Input DataFrame

required
target_schema Dict[str, str]

Target schema (column name -> type)

required
policy Any

SchemaPolicyConfig object

required

Returns:

Type Description
Any

Harmonized DataFrame

Source code in odibi/engine/base.py
@abstractmethod
def harmonize_schema(self, df: Any, target_schema: Dict[str, str], policy: Any) -> Any:
    """Harmonize DataFrame schema with target schema according to policy.

    Args:
        df: Input DataFrame
        target_schema: Target schema (column name -> type)
        policy: SchemaPolicyConfig object

    Returns:
        Harmonized DataFrame
    """
    pass

maintain_table(connection, format, table=None, path=None, config=None)

Run table maintenance operations (optimize, vacuum).

Parameters:

Name Type Description Default
connection Any

Connection object

required
format str

Table format

required
table Optional[str]

Table name

None
path Optional[str]

Table path

None
config Optional[Any]

AutoOptimizeConfig object

None
Source code in odibi/engine/base.py
def maintain_table(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    config: Optional[Any] = None,
) -> None:
    """Run table maintenance operations (optimize, vacuum).

    Args:
        connection: Connection object
        format: Table format
        table: Table name
        path: Table path
        config: AutoOptimizeConfig object
    """
    pass

materialize(df)

Materialize lazy dataset into memory (DataFrame).

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required

Returns:

Type Description
Any

Materialized DataFrame

Source code in odibi/engine/base.py
def materialize(self, df: Any) -> Any:
    """Materialize lazy dataset into memory (DataFrame).

    Args:
        df: DataFrame or LazyDataset

    Returns:
        Materialized DataFrame
    """
    return df

profile_nulls(df)

Calculate null percentage for each column.

Parameters:

Name Type Description Default
df Any

DataFrame

required

Returns:

Type Description
Dict[str, float]

Dictionary of {column_name: null_percentage} (0.0 to 1.0)

Source code in odibi/engine/base.py
def profile_nulls(self, df: Any) -> Dict[str, float]:
    """Calculate null percentage for each column.

    Args:
        df: DataFrame

    Returns:
        Dictionary of {column_name: null_percentage} (0.0 to 1.0)
    """
    return {}

read(connection, format, table=None, path=None, options=None) abstractmethod

Read data from source.

Parameters:

Name Type Description Default
connection Any

Connection object

required
format str

Data format (csv, parquet, delta, etc.)

required
table Optional[str]

Table name (for SQL/Delta)

None
path Optional[str]

File path (for file-based sources)

None
options Optional[Dict[str, Any]]

Format-specific options

None

Returns:

Type Description
Any

DataFrame (engine-specific type)

Source code in odibi/engine/base.py
@abstractmethod
def read(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    options: Optional[Dict[str, Any]] = None,
) -> Any:
    """Read data from source.

    Args:
        connection: Connection object
        format: Data format (csv, parquet, delta, etc.)
        table: Table name (for SQL/Delta)
        path: File path (for file-based sources)
        options: Format-specific options

    Returns:
        DataFrame (engine-specific type)
    """
    pass

register_format(fmt, reader=None, writer=None) classmethod

Register custom format reader/writer.

Parameters:

Name Type Description Default
fmt str

Format name (e.g. 'netcdf')

required
reader Optional[Any]

Function(path, **options) -> DataFrame

None
writer Optional[Any]

Function(df, path, **options) -> None

None
Source code in odibi/engine/base.py
@classmethod
def register_format(cls, fmt: str, reader: Optional[Any] = None, writer: Optional[Any] = None):
    """Register custom format reader/writer.

    Args:
        fmt: Format name (e.g. 'netcdf')
        reader: Function(path, **options) -> DataFrame
        writer: Function(df, path, **options) -> None
    """
    if reader:
        cls._custom_readers[fmt] = reader
    if writer:
        cls._custom_writers[fmt] = writer

table_exists(connection, table=None, path=None) abstractmethod

Check if table or location exists.

Parameters:

Name Type Description Default
connection Any

Connection object

required
table Optional[str]

Table name (for catalog tables)

None
path Optional[str]

File path (for path-based tables)

None

Returns:

Type Description
bool

True if table/location exists, False otherwise

Source code in odibi/engine/base.py
@abstractmethod
def table_exists(
    self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
) -> bool:
    """Check if table or location exists.

    Args:
        connection: Connection object
        table: Table name (for catalog tables)
        path: File path (for path-based tables)

    Returns:
        True if table/location exists, False otherwise
    """
    pass

validate_data(df, validation_config) abstractmethod

Validate data against rules.

Parameters:

Name Type Description Default
df Any

DataFrame to validate

required
validation_config Any

ValidationConfig object

required

Returns:

Type Description
List[str]

List of validation failure messages (empty if valid)

Source code in odibi/engine/base.py
@abstractmethod
def validate_data(self, df: Any, validation_config: Any) -> List[str]:
    """Validate data against rules.

    Args:
        df: DataFrame to validate
        validation_config: ValidationConfig object

    Returns:
        List of validation failure messages (empty if valid)
    """
    pass

validate_schema(df, schema_rules) abstractmethod

Validate DataFrame schema.

Parameters:

Name Type Description Default
df Any

DataFrame

required
schema_rules Dict[str, Any]

Validation rules

required

Returns:

Type Description
List[str]

List of validation failures (empty if valid)

Source code in odibi/engine/base.py
@abstractmethod
def validate_schema(self, df: Any, schema_rules: Dict[str, Any]) -> List[str]:
    """Validate DataFrame schema.

    Args:
        df: DataFrame
        schema_rules: Validation rules

    Returns:
        List of validation failures (empty if valid)
    """
    pass

write(df, connection, format, table=None, path=None, mode='overwrite', options=None, streaming_config=None) abstractmethod

Write data to destination.

Parameters:

Name Type Description Default
df Any

DataFrame to write

required
connection Any

Connection object

required
format str

Output format

required
table Optional[str]

Table name (for SQL/Delta)

None
path Optional[str]

File path (for file-based outputs)

None
mode str

Write mode (overwrite/append)

'overwrite'
options Optional[Dict[str, Any]]

Format-specific options

None
Source code in odibi/engine/base.py
@abstractmethod
def write(
    self,
    df: Any,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    mode: str = "overwrite",
    options: Optional[Dict[str, Any]] = None,
    streaming_config: Optional[Any] = None,
) -> None:
    """Write data to destination.

    Args:
        df: DataFrame to write
        connection: Connection object
        format: Output format
        table: Table name (for SQL/Delta)
        path: File path (for file-based outputs)
        mode: Write mode (overwrite/append)
        options: Format-specific options
    """
    pass

odibi.engine.pandas_engine

Pandas engine implementation.

LazyDataset dataclass

Lazy representation of a dataset (file) for out-of-core processing.

Source code in odibi/engine/pandas_engine.py
@dataclass
class LazyDataset:
    """Lazy representation of a dataset (file) for out-of-core processing."""

    path: Union[str, List[str]]
    format: str
    options: Dict[str, Any]
    connection: Optional[Any] = None  # To resolve path/credentials if needed

    def __repr__(self):
        return f"LazyDataset(path={self.path}, format={self.format})"

PandasEngine

Bases: Engine

Pandas-based execution engine.

Source code in odibi/engine/pandas_engine.py
  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
  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
 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
 433
 434
 435
 436
 437
 438
 439
 440
 441
 442
 443
 444
 445
 446
 447
 448
 449
 450
 451
 452
 453
 454
 455
 456
 457
 458
 459
 460
 461
 462
 463
 464
 465
 466
 467
 468
 469
 470
 471
 472
 473
 474
 475
 476
 477
 478
 479
 480
 481
 482
 483
 484
 485
 486
 487
 488
 489
 490
 491
 492
 493
 494
 495
 496
 497
 498
 499
 500
 501
 502
 503
 504
 505
 506
 507
 508
 509
 510
 511
 512
 513
 514
 515
 516
 517
 518
 519
 520
 521
 522
 523
 524
 525
 526
 527
 528
 529
 530
 531
 532
 533
 534
 535
 536
 537
 538
 539
 540
 541
 542
 543
 544
 545
 546
 547
 548
 549
 550
 551
 552
 553
 554
 555
 556
 557
 558
 559
 560
 561
 562
 563
 564
 565
 566
 567
 568
 569
 570
 571
 572
 573
 574
 575
 576
 577
 578
 579
 580
 581
 582
 583
 584
 585
 586
 587
 588
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
3001
3002
3003
3004
3005
3006
3007
3008
3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
3022
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
class PandasEngine(Engine):
    """Pandas-based execution engine."""

    name = "pandas"
    engine_type = EngineType.PANDAS

    def __init__(
        self,
        connections: Optional[Dict[str, Any]] = None,
        config: Optional[Dict[str, Any]] = None,
    ):
        """Initialize Pandas engine.

        Args:
            connections: Dictionary of connection objects
            config: Engine configuration (optional)
        """
        self.connections = connections or {}
        self.config = config or {}

        # Suppress noisy delta-rs transaction conflict warnings (handled by retry)
        if "RUST_LOG" not in os.environ:
            os.environ["RUST_LOG"] = "deltalake_core::kernel::transaction=error"

        # Check for performance flags
        performance = self.config.get("performance", {})

        # Determine desired state
        if hasattr(performance, "use_arrow"):
            desired_use_arrow = performance.use_arrow
        elif isinstance(performance, dict):
            desired_use_arrow = performance.get("use_arrow", True)
        else:
            desired_use_arrow = True

        # Verify availability
        if desired_use_arrow:
            try:
                import pyarrow  # noqa: F401

                self.use_arrow = True
            except ImportError:
                import logging

                logger = logging.getLogger(__name__)
                logger.warning(
                    "Apache Arrow not found. Disabling Arrow optimizations. "
                    "Install 'pyarrow' to enable."
                )
                self.use_arrow = False
        else:
            self.use_arrow = False

        # Check for DuckDB
        self.use_duckdb = False
        # Default to False to ensure stability with existing tests (Lazy Loading is opt-in)
        if self.config.get("performance", {}).get("use_duckdb", False):
            try:
                import duckdb  # noqa: F401

                self.use_duckdb = True
            except ImportError:
                pass

    def materialize(self, df: Any) -> Any:
        """Materialize lazy dataset.

        Args:
            df: DataFrame or LazyDataset to materialize.

        Returns:
            Materialized pandas DataFrame. If input is already a DataFrame,
            returns it unchanged.
        """
        if isinstance(df, LazyDataset):
            # Re-invoke read but force materialization (by bypassing Lazy check)
            # We pass the resolved path directly
            # Note: We need to handle the case where path was resolved.
            # LazyDataset.path should be the FULL path.
            return self._read_file(
                full_path=df.path, format=df.format, options=df.options, connection=df.connection
            )
        return df

    def _process_df(
        self, df: Union[pd.DataFrame, Iterator[pd.DataFrame]], query: Optional[str]
    ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
        """Apply post-read processing (filtering)."""
        if query and df is not None:
            # Handle Iterator
            from collections.abc import Iterator

            if isinstance(df, Iterator):
                # Filter each chunk
                return (chunk.query(query) for chunk in df)

            if not df.empty:
                try:
                    return df.query(query)
                except Exception as e:
                    import logging

                    logger = logging.getLogger(__name__)
                    logger.warning(f"Failed to apply query '{query}': {e}")
        return df

    _CLOUD_URI_PREFIXES = (
        "abfss://",
        "abfs://",
        "wasbs://",
        "wasb://",
        "az://",
        "s3://",
        "gs://",
        "https://",
    )

    def _retry_delta_operation(self, func, max_retries: int = 5, base_delay: float = 0.2):
        """Retry Delta operations with exponential backoff for concurrent conflicts."""
        for attempt in range(max_retries):
            try:
                return func()
            except Exception as e:
                error_str = str(e).lower()
                is_conflict = "conflict" in error_str or "concurrent" in error_str
                if attempt == max_retries - 1 or not is_conflict:
                    raise
                delay = base_delay * (2**attempt) + random.uniform(0, 0.1)
                time.sleep(delay)

    def _read_delta_with_duckdb(
        self,
        path: str,
        storage_options: dict,
        version: Optional[int],
        timestamp: Optional[str],
        post_read_query: Optional[str],
        ctx,
    ) -> pd.DataFrame:
        """Read Delta table using DuckDB (supports DeletionVectors and ColumnMapping).

        DuckDB's delta extension supports advanced Delta features that the Python
        deltalake library doesn't yet support. This is used as a fallback when
        deltalake fails with unsupported reader features.

        Args:
            path: Path to Delta table
            storage_options: Storage authentication options (for ADLS, S3, etc.)
            version: Specific version to read (time travel)
            timestamp: Specific timestamp to read (time travel)
            post_read_query: Optional pandas query to apply after reading
            ctx: Logging context

        Returns:
            DataFrame with Delta table contents
        """
        try:
            import duckdb
        except ImportError:
            raise ImportError(
                "DuckDB is required to read Delta tables with DeletionVectors or ColumnMapping. "
                "Install with 'pip install duckdb>=1.0.0'."
            )

        ctx.debug("Reading Delta table with DuckDB", path=path)

        conn = duckdb.connect(":memory:")

        # Install and load the delta extension
        conn.execute("INSTALL delta")
        conn.execute("LOAD delta")

        # Configure storage options for cloud paths
        if storage_options:
            # Azure ADLS configuration
            if "account_name" in storage_options:
                account = storage_options["account_name"]
                if "account_key" in storage_options:
                    conn.execute(
                        f"SET azure_storage_account_name = '{account}';"
                        f"SET azure_storage_account_key = '{storage_options['account_key']}';"
                    )
                elif "sas_token" in storage_options:
                    sas = storage_options["sas_token"]
                    if sas.startswith("?"):
                        sas = sas[1:]
                    conn.execute(
                        f"SET azure_storage_account_name = '{account}';"
                        f"SET azure_storage_sas_token = '{sas}';"
                    )
                # For Azure Identity (DefaultCredential), DuckDB uses env vars automatically

            # AWS S3 configuration
            if "AWS_ACCESS_KEY_ID" in storage_options:
                conn.execute(
                    f"SET s3_access_key_id = '{storage_options['AWS_ACCESS_KEY_ID']}';"
                    f"SET s3_secret_access_key = '{storage_options['AWS_SECRET_ACCESS_KEY']}';"
                )
                if "AWS_REGION" in storage_options:
                    conn.execute(f"SET s3_region = '{storage_options['AWS_REGION']}';")

        # Build the delta_scan query
        # DuckDB delta_scan doesn't support version/timestamp time travel
        if version is not None or timestamp is not None:
            raise NotImplementedError(
                "DuckDB delta_scan does not support time travel. "
                "Install 'deltalake' package for time travel support, "
                "or remove the version/timestamp option."
            )

        query = f"SELECT * FROM delta_scan('{path}')"
        ctx.debug("Executing DuckDB query", query=query)

        df = conn.execute(query).fetchdf()
        conn.close()

        ctx.debug("Delta table read via DuckDB", rows=len(df), columns=len(df.columns))

        return self._process_df(df, post_read_query)

    def _resolve_path(self, path: Optional[str], connection: Any) -> str:
        """Resolve path to full URI, avoiding double-prefixing for cloud URIs.

        Args:
            path: Relative or absolute path
            connection: Connection object (may have get_path method)

        Returns:
            Full resolved path
        """
        if not path:
            raise ValueError(
                "Failed to resolve path: path argument is required but was empty or None. "
                "Provide a valid file path or use 'table' parameter with a connection."
            )
        if path.startswith(self._CLOUD_URI_PREFIXES):
            return path
        if connection:
            return connection.get_path(path)
        return path

    def _merge_storage_options(
        self, connection: Any, options: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Merge connection storage options with user options.

        Args:
            connection: Connection object (may have pandas_storage_options method)
            options: User-provided options

        Returns:
            Merged options dictionary
        """
        options = options or {}

        # If connection provides storage_options (e.g., AzureADLS), merge them
        if hasattr(connection, "pandas_storage_options"):
            conn_storage_opts = connection.pandas_storage_options()
            user_storage_opts = options.get("storage_options", {})

            # User options override connection options
            merged_storage_opts = {**conn_storage_opts, **user_storage_opts}

            # Return options with merged storage_options
            return {**options, "storage_options": merged_storage_opts}

        return options

    def _is_remote_uri(self, path: Union[str, Path]) -> bool:
        """Check if a path is a remote URI (cloud storage).

        Args:
            path: File path or URI

        Returns:
            True if path is a remote URI requiring fsspec
        """
        if isinstance(path, Path):
            path = str(path)

        parsed = urlparse(path)

        # Windows drive letter check (e.g., C:\foo.xlsx parses as scheme='c')
        if os.name == "nt" and len(parsed.scheme) == 1 and parsed.scheme.isalpha():
            return False

        # Cloud schemes
        if parsed.scheme and parsed.scheme not in ("file", ""):
            return True

        # Fallback prefix check
        return path.startswith(self._CLOUD_URI_PREFIXES)

    def _expand_remote_glob(
        self,
        pattern: str,
        storage_options: Optional[Dict[str, Any]] = None,
        ctx: Any = None,
    ) -> List[str]:
        """Expand glob pattern for remote paths using fsspec.

        Args:
            pattern: Glob pattern (e.g., "abfss://container@account.dfs.core.windows.net/path/*.xlsx")
            storage_options: Authentication options for cloud storage
            ctx: Logging context

        Returns:
            List of matching file URIs
        """
        if ctx is None:
            ctx = get_logging_context()

        try:
            import fsspec
        except ImportError:
            ctx.error("fsspec required for remote glob expansion")
            raise ImportError(
                "Remote glob patterns require 'fsspec'. Install with: pip install fsspec adlfs"
            )

        # Parse protocol and path
        if "://" in pattern:
            protocol, _, path_part = pattern.partition("://")
        else:
            protocol = "file"
            path_part = pattern

        ctx.debug(
            "Expanding remote glob pattern",
            protocol=protocol,
            pattern=path_part,
        )

        # Create filesystem with storage options
        fs = fsspec.filesystem(protocol, **(storage_options or {}))

        # Glob and reconstruct full URIs
        matched = fs.glob(path_part)

        if not matched:
            ctx.warning("No files matched remote glob pattern", pattern=pattern)
            return []

        # Reconstruct full URIs
        result = [f"{protocol}://{m}" for m in matched]

        ctx.info(
            "Remote glob pattern expanded",
            pattern=pattern,
            matched_files=len(result),
        )

        return result

    def _open_remote_file(
        self,
        path: str,
        storage_options: Optional[Dict[str, Any]] = None,
    ):
        """Open a remote file using fsspec.

        Args:
            path: Remote file URI
            storage_options: Authentication options for cloud storage

        Returns:
            Context manager yielding a binary file-like object
        """
        import fsspec

        return fsspec.open(path, "rb", **(storage_options or {}))

    def _read_parallel(self, read_func: Any, paths: List[str], **kwargs) -> pd.DataFrame:
        """Read multiple files in parallel using threads.

        Args:
            read_func: Pandas read function (e.g. pd.read_csv)
            paths: List of file paths
            kwargs: Arguments to pass to read_func

        Returns:
            Concatenated DataFrame
        """
        # Conservative worker count to avoid OOM on large files
        max_workers = min(8, os.cpu_count() or 4)

        dfs = []
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # map preserves order
            results = executor.map(lambda p: read_func(p, **kwargs), paths)
            dfs = list(results)

        non_empty = [df for df in dfs if not df.empty]
        return pd.concat(non_empty, ignore_index=True) if non_empty else pd.DataFrame()

    def _read_excel_with_patterns(
        self,
        paths: Union[str, Path, List[str]],
        sheet_pattern: Union[str, List[str], None] = None,
        sheet_pattern_case_sensitive: bool = False,
        add_source_file: bool = False,
        is_glob: bool = False,
        ctx: Any = None,
        storage_options: Optional[Dict[str, Any]] = None,
        **read_kwargs,
    ) -> pd.DataFrame:
        """Read Excel files with glob and sheet pattern support.

        Supports both local files and remote cloud storage (Azure Blob, S3, etc.)
        via fsspec integration.

        Args:
            paths: Single path, glob-expanded list, or list of paths.
                   Can be local paths or cloud URIs (abfss://, s3://, etc.)
            sheet_pattern: Pattern(s) to match sheet names (e.g., "*powerbi*").
                          Supports single string or list of patterns.
            sheet_pattern_case_sensitive: Whether pattern matching is case-sensitive.
            add_source_file: If True, adds '_source_file' column with filename.
            is_glob: Whether paths contains a glob pattern to expand.
            ctx: Logging context.
            storage_options: Authentication options for cloud storage (from connection).
            **read_kwargs: Additional kwargs for pd.read_excel.

        Returns:
            Concatenated DataFrame from all matching files/sheets.
        """
        if ctx is None:
            ctx = get_logging_context()

        # Normalize patterns to list
        if sheet_pattern is None:
            patterns = None
        elif isinstance(sheet_pattern, str):
            patterns = [sheet_pattern]
        else:
            patterns = list(sheet_pattern)

        def match_sheet(sheet_name: str) -> bool:
            """Check if sheet name matches any pattern."""
            if patterns is None:
                return True  # No pattern = match all

            check_name = sheet_name if sheet_pattern_case_sensitive else sheet_name.lower()
            for pattern in patterns:
                check_pattern = pattern if sheet_pattern_case_sensitive else pattern.lower()
                if fnmatch.fnmatch(check_name, check_pattern):
                    return True
            return False

        def get_file_name(file_path: str) -> str:
            """Extract filename from path (works for both local and remote URIs)."""
            if "://" in file_path:
                # Remote URI - extract filename from path part
                path_part = file_path.split("://", 1)[1]
                return path_part.rsplit("/", 1)[-1] if "/" in path_part else path_part
            else:
                return Path(file_path).name

        def read_single_excel(file_path: str) -> pd.DataFrame:
            """Read a single Excel file, matching sheets by pattern.

            Handles both local and remote files transparently.
            """
            file_name = get_file_name(file_path)
            is_remote = self._is_remote_uri(file_path)

            # For remote files, remove storage_options - we handle it via fsspec
            # For local files, keep storage_options for consistency (pandas ignores it)
            if is_remote:
                excel_kwargs = {k: v for k, v in read_kwargs.items() if k != "storage_options"}
            else:
                excel_kwargs = read_kwargs.copy()
                # Add storage_options back if provided (for test compatibility)
                if storage_options and "storage_options" not in excel_kwargs:
                    excel_kwargs["storage_options"] = storage_options

            try:
                if is_remote:
                    # Remote file - use fsspec to open and pass file handle
                    ctx.debug(
                        "Reading remote Excel file",
                        file=file_name,
                        uri=file_path,
                    )
                    with self._open_remote_file(file_path, storage_options) as fh:
                        return _read_excel_from_source(fh, file_name, excel_kwargs)
                else:
                    # Local file - pass path directly to pandas (supports mocking)
                    ctx.debug("Reading local Excel file", file=file_name)
                    return _read_excel_from_source(file_path, file_name, excel_kwargs)

            except Exception as e:
                ctx.warning(
                    "Failed to read Excel file",
                    file=file_name,
                    path=file_path,
                    error=str(e),
                )
                raise

        def _read_excel_from_source(source, file_name: str, excel_kwargs: dict) -> pd.DataFrame:
            """Read Excel data from a path or file handle with sheet pattern matching.

            Uses fastexcel (Rust/calamine) when available for 5-10x faster Excel
            parsing compared to openpyxl. Falls back to pandas/openpyxl if
            fastexcel is not installed.

            Args:
                source: File path (str) or file handle for pandas to read from
                file_name: Display name for the file (used in _source_file column)
                excel_kwargs: Keyword arguments to pass to pd.read_excel
            """
            # Try fastexcel (Rust/calamine) first — much faster than openpyxl
            try:
                import fastexcel

                return _read_excel_fastexcel(source, file_name, excel_kwargs, fastexcel)
            except ImportError:
                pass
            except Exception as e:
                ctx.debug(
                    "fastexcel failed, falling back to openpyxl",
                    file=file_name,
                    error=str(e),
                )

            # Fallback to pandas/openpyxl
            return _read_excel_pandas(source, file_name, excel_kwargs)

        def _read_excel_fastexcel(
            source, file_name: str, excel_kwargs: dict, fastexcel
        ) -> pd.DataFrame:
            """Read Excel using fastexcel/calamine (Rust-based, fast)."""
            import logging

            # fastexcel needs bytes for file handles
            if hasattr(source, "read"):
                source_data = source.read()
                if hasattr(source, "seek"):
                    source.seek(0)
            else:
                source_data = source

            parser = fastexcel.read_excel(source_data)
            all_sheet_names = parser.sheet_names

            # Determine which sheets to read
            if patterns is None:
                # No pattern — read first sheet only
                sheet_name_hint = excel_kwargs.get("sheet_name", 0)
                if isinstance(sheet_name_hint, int):
                    sheets_to_read = [all_sheet_names[sheet_name_hint]]
                elif isinstance(sheet_name_hint, str):
                    sheets_to_read = [sheet_name_hint]
                else:
                    sheets_to_read = [all_sheet_names[0]]
            else:
                sheets_to_read = [s for s in all_sheet_names if match_sheet(s)]

            if not sheets_to_read:
                ctx.debug(
                    "No matching sheets in Excel file",
                    file=file_name,
                    patterns=patterns,
                    available_sheets=all_sheet_names,
                )
                return pd.DataFrame()

            ctx.debug(
                "Reading Excel sheets (fastexcel/calamine)",
                file=file_name,
                matching_sheets=sheets_to_read,
            )

            # Check if dtype=str requested (from YAML it arrives as string "str")
            force_dtype = excel_kwargs.get("dtype")
            force_string = force_dtype is str or force_dtype == "str"

            # Suppress "Could not determine dtype" warnings from fastexcel.
            # These come through Python's logging module as logger
            # "fastexcel.types.dtype" at WARNING level.
            fe_logger = logging.getLogger("fastexcel.types.dtype")
            original_level = fe_logger.level
            fe_logger.setLevel(logging.ERROR)

            dfs = []
            try:
                for sheet in sheets_to_read:
                    ws = parser.load_sheet_by_name(sheet)
                    df = ws.to_pandas()
                    if force_string:
                        df = df.astype(str)
                        # astype(str) turns NaN/NaT/None into literal strings
                        # "nan", "NaT", "None" — restore them to real nulls
                        # so SQL IS NOT NULL filters work as expected.
                        df = df.replace({"nan": None, "NaT": None, "None": None})
                    if add_source_file:
                        df["_source_file"] = file_name
                        df["_source_sheet"] = sheet
                    dfs.append(df)
            finally:
                fe_logger.setLevel(original_level)

            non_empty = [df for df in dfs if not df.empty]
            return pd.concat(non_empty, ignore_index=True) if non_empty else pd.DataFrame()

        def _read_excel_pandas(source, file_name: str, excel_kwargs: dict) -> pd.DataFrame:
            """Read Excel using pandas/openpyxl (slower fallback)."""
            force_dtype = excel_kwargs.get("dtype")
            _force_str = force_dtype is str or force_dtype == "str"

            def _fix_str_nulls(df: pd.DataFrame) -> pd.DataFrame:
                """Replace stringified nulls back to real None after dtype=str."""
                if _force_str:
                    df = df.replace({"nan": None, "NaT": None, "None": None})
                return df

            # If no sheet pattern, use simple read_excel
            if patterns is None:
                df = _fix_str_nulls(pd.read_excel(source, **excel_kwargs))
                if add_source_file:
                    df["_source_file"] = file_name
                    df["_source_sheet"] = excel_kwargs.get("sheet_name", 0)
                return df

            # With sheet pattern, need to inspect sheets first
            xls = pd.ExcelFile(source)
            matching_sheets = [s for s in xls.sheet_names if match_sheet(s)]

            if not matching_sheets:
                ctx.debug(
                    "No matching sheets in Excel file",
                    file=file_name,
                    patterns=patterns,
                    available_sheets=xls.sheet_names,
                )
                return pd.DataFrame()

            ctx.debug(
                "Reading Excel sheets (openpyxl)",
                file=file_name,
                matching_sheets=matching_sheets,
            )

            dfs = []
            for sheet in matching_sheets:
                df = _fix_str_nulls(pd.read_excel(xls, sheet_name=sheet, **excel_kwargs))
                if add_source_file:
                    df["_source_file"] = file_name
                    df["_source_sheet"] = sheet
                dfs.append(df)

            non_empty = [df for df in dfs if not df.empty]
            return pd.concat(non_empty, ignore_index=True) if non_empty else pd.DataFrame()

        # Expand glob patterns if needed
        file_list: List[str] = []

        if isinstance(paths, (str, Path)):
            path_str = str(paths)

            # Check if this is a glob pattern that needs expansion
            has_glob = "*" in path_str or "?" in path_str or "[" in path_str

            if has_glob:
                if self._is_remote_uri(path_str):
                    # Remote glob - use fsspec
                    file_list = self._expand_remote_glob(path_str, storage_options, ctx)
                else:
                    # Local glob
                    file_list = glob.glob(path_str)
                    if not file_list:
                        raise FileNotFoundError(f"No files matched pattern: {path_str}")
                    ctx.info(
                        "Local glob pattern expanded",
                        pattern=path_str,
                        matched_files=len(file_list),
                    )
            else:
                # Single file (no glob)
                file_list = [path_str]

        elif isinstance(paths, list):
            # Already a list of paths (pre-expanded)
            file_list = [str(p) for p in paths]

        if not file_list:
            raise FileNotFoundError(f"No Excel files found for: {paths}")

        # Read all files
        ctx.info(
            "Reading Excel files",
            file_count=len(file_list),
            sheet_patterns=patterns,
            is_remote=self._is_remote_uri(file_list[0]) if file_list else False,
        )

        all_dfs = []
        source_files = []

        for file_path in file_list:
            df = read_single_excel(file_path)
            if not df.empty:
                all_dfs.append(df)
                source_files.append(file_path)

        if not all_dfs:
            ctx.warning("No data read from Excel files", file_count=len(file_list))
            return pd.DataFrame()

        result = pd.concat(all_dfs, ignore_index=True)
        if hasattr(result, "attrs"):
            result.attrs["odibi_source_files"] = source_files

        return result

    def _read_simulation(self, options: Dict[str, Any], ctx: Any) -> pd.DataFrame:
        """Read simulated data.

        Args:
            options: Read options containing simulation configuration
            ctx: Logging context

        Returns:
            DataFrame with simulated data
        """
        from odibi.config import SimulationConfig
        from odibi.simulation import SimulationEngine

        # Extract simulation config
        sim_config_dict = options.get("simulation")
        if not sim_config_dict:
            raise ValueError(
                "Simulation format requires 'simulation' key in read options. "
                "Example: read.options.simulation = {scope: {...}, entities: {...}, columns: [...]}"
            )

        # Parse and validate simulation config
        sim_config = SimulationConfig(**sim_config_dict)

        # Validate write compatibility if write info passed
        write_mode = options.get("_write_mode")
        write_keys = options.get("_write_keys")
        if write_mode:
            sim_config.validate_write_compatibility(write_mode, write_keys)

        # Check for HWM (passed through options for incremental mode)
        hwm_timestamp = options.get("_hwm_timestamp")

        # Check for random walk state (passed through options for incremental mode)
        random_walk_state = options.get("_random_walk_state")

        # Check for scheduled event state (passed through options for incremental mode)
        scheduled_event_state = options.get("_scheduled_event_state")

        # Check for stateful-function state (prev/ema/pid/delay) for incremental continuity
        entity_state = options.get("_entity_state")

        # Create simulation engine
        engine = SimulationEngine(
            sim_config,
            hwm_timestamp=hwm_timestamp,
            random_walk_state=random_walk_state,
            scheduled_event_state=scheduled_event_state,
            entity_state=entity_state,
        )

        # Generate data
        ctx.info(
            "Generating simulated data",
            entities=len(engine.entity_names),
            rows_per_entity=engine.total_rows,
            total_rows=len(engine.entity_names) * engine.total_rows,
        )

        rows = engine.generate()

        # Convert to DataFrame
        df = pd.DataFrame(rows)

        # Cast int columns to nullable Int64 (avoids float promotion when nulls exist)
        # Use Float64 as intermediate step to handle object columns with mixed
        # None/float values (e.g., from derived expressions using round(..., 0))
        for col_config in sim_config.columns:
            if col_config.data_type == "int" and col_config.name in df.columns:
                df[col_config.name] = df[col_config.name].astype("Float64").round().astype("Int64")

        # Store max timestamp for HWM tracking
        max_ts = engine.get_max_timestamp(rows)
        if max_ts and hasattr(df, "attrs"):
            df.attrs["_simulation_max_timestamp"] = max_ts

        # Store random walk final state for incremental persistence
        rw_state = engine.get_random_walk_final_state(rows)
        if rw_state and hasattr(df, "attrs"):
            df.attrs["_simulation_random_walk_state"] = rw_state

        # Store scheduled event final state for incremental persistence
        se_state = engine.get_scheduled_event_final_state()
        if se_state and hasattr(df, "attrs"):
            df.attrs["_simulation_scheduled_event_state"] = se_state

        # Store stateful-function state (prev/ema/pid/delay) for incremental persistence
        entity_state_final = engine.get_entity_state_final()
        if entity_state_final and hasattr(df, "attrs"):
            df.attrs["_simulation_entity_state"] = entity_state_final

        ctx.info(
            "Simulation complete",
            rows=len(df),
            columns=len(df.columns) if not df.empty else 0,
        )

        return df

    def read(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        streaming: bool = False,
        schema: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        as_of_version: Optional[int] = None,
        as_of_timestamp: Optional[str] = None,
    ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
        """Read data using Pandas (or LazyDataset).

        Args:
            connection: Connection object providing base path and storage options.
            format: Data format (csv, parquet, delta, json, excel, avro, sql).
            table: Table name (mutually exclusive with path).
            path: File path (mutually exclusive with table).
            streaming: Whether to enable streaming mode (not supported in Pandas).
            schema: Optional schema specification (not used in Pandas engine).
            options: Additional read options to pass to pandas readers.
            as_of_version: Delta Lake version for time travel queries.
            as_of_timestamp: Delta Lake timestamp for time travel queries.

        Returns:
            DataFrame or iterator of DataFrames containing the loaded data.

        Raises:
            ValueError: If streaming is requested or neither path nor table is provided.
        """
        ctx = get_logging_context().with_context(engine="pandas")
        start = time.time()

        source = path or table
        ctx.debug(
            "Starting read operation",
            format=format,
            path=source,
            streaming=streaming,
            use_arrow=self.use_arrow,
        )

        if streaming:
            ctx.error(
                "Streaming not supported in Pandas engine",
                format=format,
                path=source,
            )
            raise ValueError(
                "Streaming is not supported in the Pandas engine. "
                "Please use 'engine: spark' for streaming pipelines."
            )

        options = options or {}

        # Handle simulation format (no path/table needed)
        if format == "simulation":
            ctx.debug("Generating simulated data")
            return self._read_simulation(options, ctx)

        # Resolve full path from connection
        try:
            full_path = self._resolve_path(path or table, connection)
        except ValueError:
            if table and not connection:
                ctx.error("Connection required when specifying 'table'", table=table)
                raise ValueError(
                    f"Cannot read table '{table}': connection is required when using 'table' parameter. "
                    "Provide a valid connection object or use 'path' for file-based reads."
                )
            ctx.error("Neither path nor table provided for read operation")
            raise ValueError(
                "Read operation failed: neither 'path' nor 'table' was provided. "
                "Specify a file path or table name in your configuration."
            )

        # Merge storage options for cloud connections
        merged_options = self._merge_storage_options(connection, options)

        # Sanitize options for pandas compatibility
        if "header" in merged_options:
            if merged_options["header"] is True:
                merged_options["header"] = 0
            elif merged_options["header"] is False:
                merged_options["header"] = None

        # Handle Time Travel options
        if as_of_version is not None:
            merged_options["versionAsOf"] = as_of_version
            ctx.debug("Time travel enabled", version=as_of_version)
        if as_of_timestamp is not None:
            merged_options["timestampAsOf"] = as_of_timestamp
            ctx.debug("Time travel enabled", timestamp=as_of_timestamp)

        # Check for Lazy/DuckDB optimization
        can_lazy_load = False

        if can_lazy_load:
            ctx.debug("Using lazy loading via DuckDB", path=str(full_path))
            if isinstance(full_path, (str, Path)):
                return LazyDataset(
                    path=str(full_path),
                    format=format,
                    options=merged_options,
                    connection=connection,
                )
            elif isinstance(full_path, list):
                return LazyDataset(
                    path=full_path, format=format, options=merged_options, connection=connection
                )

        result = self._read_file(full_path, format, merged_options, connection)

        # Log metrics for materialized DataFrames
        elapsed = (time.time() - start) * 1000
        if isinstance(result, pd.DataFrame):
            row_count = len(result)
            memory_mb = result.memory_usage(deep=True).sum() / (1024 * 1024)

            ctx.log_file_io(
                path=str(full_path) if not isinstance(full_path, list) else str(full_path[0]),
                format=format,
                mode="read",
                rows=row_count,
            )
            ctx.log_pandas_metrics(
                memory_mb=memory_mb,
                dtypes={col: str(dtype) for col, dtype in result.dtypes.items()},
            )
            ctx.info(
                "Read completed",
                format=format,
                rows=row_count,
                elapsed_ms=round(elapsed, 2),
                memory_mb=round(memory_mb, 2),
            )

        return result

    def _read_file(
        self,
        full_path: Union[str, List[str], Any],
        format: str,
        options: Dict[str, Any],
        connection: Any = None,
    ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
        """Internal file reading logic."""
        ctx = get_logging_context().with_context(engine="pandas")

        ctx.debug(
            "Reading file",
            path=str(full_path) if not isinstance(full_path, list) else f"{len(full_path)} files",
            format=format,
        )

        # Custom Readers
        if format in self._custom_readers:
            ctx.debug(f"Using custom reader for format: {format}")
            return self._custom_readers[format](full_path, **options)

        # Handle glob patterns for local files
        is_glob = False
        if isinstance(full_path, (str, Path)) and (
            "*" in str(full_path) or "?" in str(full_path) or "[" in str(full_path)
        ):
            parsed = urlparse(str(full_path))
            # Only expand for local files (no scheme, file://, or drive letter)
            is_local = (
                not parsed.scheme
                or parsed.scheme == "file"
                or (len(parsed.scheme) == 1 and parsed.scheme.isalpha())
            )

            if is_local:
                glob_path = str(full_path)
                if glob_path.startswith("file:///"):
                    glob_path = glob_path[8:]
                elif glob_path.startswith("file://"):
                    glob_path = glob_path[7:]

                matched_files = glob.glob(glob_path)
                if not matched_files:
                    ctx.error(
                        "No files matched glob pattern",
                        pattern=glob_path,
                    )
                    raise FileNotFoundError(f"No files matched pattern: {glob_path}")

                ctx.info(
                    "Glob pattern expanded",
                    pattern=glob_path,
                    matched_files=len(matched_files),
                )
                full_path = matched_files
                is_glob = True

        # Prepare read options (options already includes storage_options from caller)
        read_kwargs = options.copy()

        # Filter out Spark-specific options that don't apply to Pandas
        spark_only_options = {
            "inferSchema",
            "multiLine",
            "mode",
            "columnNameOfCorruptRecord",
            "dateFormat",
            "timestampFormat",
            "nullValue",
            "nanValue",
            "positiveInf",
            "negativeInf",
            "escape",
            "charToEscapeQuoteEscaping",
            "ignoreLeadingWhiteSpace",
            "ignoreTrailingWhiteSpace",
            "maxColumns",
            "maxCharsPerColumn",
            "unescapedQuoteHandling",
            "enforceSchema",
            "samplingRatio",
            "emptyValue",
            "locale",
            "lineSep",
            "pathGlobFilter",
            "recursiveFileLookup",
            "modifiedBefore",
            "modifiedAfter",
        }
        for opt in spark_only_options:
            read_kwargs.pop(opt, None)

        # Extract 'query' or 'filter' option for post-read filtering
        post_read_query = read_kwargs.pop("query", None) or read_kwargs.pop("filter", None)

        if self.use_arrow:
            read_kwargs["dtype_backend"] = "pyarrow"

        # Read based on format
        if format == "csv":
            try:
                if is_glob and isinstance(full_path, list):
                    ctx.debug(
                        "Parallel CSV read",
                        file_count=len(full_path),
                    )
                    df = self._read_parallel(pd.read_csv, full_path, **read_kwargs)
                    df.attrs["odibi_source_files"] = full_path
                    return self._process_df(df, post_read_query)

                df = pd.read_csv(full_path, **read_kwargs)
                if hasattr(df, "attrs"):
                    df.attrs["odibi_source_files"] = [str(full_path)]
                return self._process_df(df, post_read_query)
            except UnicodeDecodeError:
                ctx.warning(
                    "UnicodeDecodeError, retrying with latin1 encoding",
                    path=str(full_path),
                )
                read_kwargs["encoding"] = "latin1"
                try:
                    if is_glob and isinstance(full_path, list):
                        df = self._read_parallel(pd.read_csv, full_path, **read_kwargs)
                        df.attrs["odibi_source_files"] = full_path
                        return self._process_df(df, post_read_query)

                    df = pd.read_csv(full_path, **read_kwargs)
                    if hasattr(df, "attrs"):
                        df.attrs["odibi_source_files"] = [str(full_path)]
                    return self._process_df(df, post_read_query)
                except pd.errors.ParserError:
                    ctx.warning(
                        "ParserError after encoding fix, retrying with on_bad_lines='skip'",
                        path=str(full_path),
                    )
                    read_kwargs["on_bad_lines"] = "skip"
                    if is_glob and isinstance(full_path, list):
                        df = self._read_parallel(pd.read_csv, full_path, **read_kwargs)
                        df.attrs["odibi_source_files"] = full_path
                        return self._process_df(df, post_read_query)

                    df = pd.read_csv(full_path, **read_kwargs)
                    if hasattr(df, "attrs"):
                        df.attrs["odibi_source_files"] = [str(full_path)]
                    return self._process_df(df, post_read_query)
            except pd.errors.ParserError:
                ctx.warning(
                    "ParserError, retrying with on_bad_lines='skip'",
                    path=str(full_path),
                )
                read_kwargs["on_bad_lines"] = "skip"
                if is_glob and isinstance(full_path, list):
                    df = self._read_parallel(pd.read_csv, full_path, **read_kwargs)
                    df.attrs["odibi_source_files"] = full_path
                    return self._process_df(df, post_read_query)

                df = pd.read_csv(full_path, **read_kwargs)
                if hasattr(df, "attrs"):
                    df.attrs["odibi_source_files"] = [str(full_path)]
                return self._process_df(df, post_read_query)
        elif format == "parquet":
            ctx.debug("Reading parquet", path=str(full_path))
            df = pd.read_parquet(full_path, **read_kwargs)
            if isinstance(full_path, list):
                df.attrs["odibi_source_files"] = full_path
            else:
                df.attrs["odibi_source_files"] = [str(full_path)]
            return self._process_df(df, post_read_query)
        elif format == "json":
            if is_glob and isinstance(full_path, list):
                ctx.debug(
                    "Parallel JSON read",
                    file_count=len(full_path),
                )
                df = self._read_parallel(pd.read_json, full_path, **read_kwargs)
                df.attrs["odibi_source_files"] = full_path
                return self._process_df(df, post_read_query)

            df = pd.read_json(full_path, **read_kwargs)
            if hasattr(df, "attrs"):
                df.attrs["odibi_source_files"] = [str(full_path)]
            return self._process_df(df, post_read_query)
        elif format == "excel":
            ctx.debug("Reading Excel file", path=str(full_path))
            read_kwargs.pop("dtype_backend", None)

            # Extract excel-specific options
            sheet_pattern = read_kwargs.pop("sheet_pattern", None)
            sheet_pattern_case_sensitive = read_kwargs.pop("sheet_pattern_case_sensitive", False)
            add_source_file = read_kwargs.pop("add_source_file", False)

            # Extract storage_options for cloud storage authentication
            storage_options = read_kwargs.pop("storage_options", None)

            df = self._read_excel_with_patterns(
                full_path,
                sheet_pattern=sheet_pattern,
                sheet_pattern_case_sensitive=sheet_pattern_case_sensitive,
                add_source_file=add_source_file,
                is_glob=is_glob,
                ctx=ctx,
                storage_options=storage_options,
                **read_kwargs,
            )
            return self._process_df(df, post_read_query)
        elif format == "delta":
            ctx.debug("Reading Delta table", path=str(full_path))
            try:
                from deltalake import DeltaTable
            except ImportError:
                ctx.error(
                    "Delta Lake library not installed",
                    path=str(full_path),
                )
                raise ImportError(
                    "Delta Lake support requires 'pip install odibi[pandas]' "
                    "or 'pip install deltalake'. See README.md for installation instructions."
                )

            storage_opts = options.get("storage_options", {})
            version = options.get("versionAsOf")
            timestamp = options.get("timestampAsOf")

            # Try deltalake first, fall back to DuckDB for unsupported features
            # (e.g., DeletionVectors, ColumnMapping)
            try:
                if timestamp is not None:
                    from datetime import datetime as dt_module

                    if isinstance(timestamp, str):
                        ts = dt_module.fromisoformat(timestamp.replace("Z", "+00:00"))
                    else:
                        ts = timestamp
                    dt = DeltaTable(full_path, storage_options=storage_opts)
                    dt.load_with_datetime(ts)
                    ctx.debug("Delta table loaded with timestamp", timestamp=str(ts))
                elif version is not None:
                    dt = DeltaTable(full_path, storage_options=storage_opts, version=version)
                    ctx.debug("Delta table loaded with version", version=version)
                else:
                    dt = DeltaTable(full_path, storage_options=storage_opts)
                    ctx.debug("Delta table loaded (latest version)")

                if self.use_arrow:
                    import inspect

                    sig = inspect.signature(dt.to_pandas)

                    if "arrow_options" in sig.parameters:
                        return self._process_df(
                            dt.to_pandas(
                                partitions=None, arrow_options={"types_mapper": pd.ArrowDtype}
                            ),
                            post_read_query,
                        )
                    else:
                        return self._process_df(
                            dt.to_pyarrow_table().to_pandas(types_mapper=pd.ArrowDtype),
                            post_read_query,
                        )
                else:
                    return self._process_df(dt.to_pandas(), post_read_query)

            except Exception as e:
                # Check if this is a DeltaProtocolError for unsupported features
                error_msg = str(e).lower()
                unsupported_features = (
                    "deletionvectors" in error_msg
                    or "columnmapping" in error_msg
                    or "reader features" in error_msg
                    or "not yet supported" in error_msg
                )

                if unsupported_features:
                    ctx.warning(
                        "Delta table uses features not supported by deltalake library, "
                        "falling back to DuckDB",
                        path=str(full_path),
                        error=str(e),
                    )
                    return self._read_delta_with_duckdb(
                        full_path, storage_opts, version, timestamp, post_read_query, ctx
                    )
                else:
                    raise
        elif format == "avro":
            ctx.debug("Reading Avro file", path=str(full_path))
            try:
                import fastavro
            except ImportError:
                ctx.error(
                    "fastavro library not installed",
                    path=str(full_path),
                )
                raise ImportError(
                    "Avro support requires 'pip install odibi[pandas]' "
                    "or 'pip install fastavro'. See README.md for installation instructions."
                )

            parsed = urlparse(full_path)
            if parsed.scheme and parsed.scheme not in ["file", ""]:
                import fsspec

                storage_opts = options.get("storage_options", {})
                with fsspec.open(full_path, "rb", **storage_opts) as f:
                    reader = fastavro.reader(f)
                    records = [record for record in reader]
                return pd.DataFrame(records)
            else:
                with open(full_path, "rb") as f:
                    reader = fastavro.reader(f)
                    records = [record for record in reader]
                return self._process_df(pd.DataFrame(records), post_read_query)
        elif is_sql_format(format):
            ctx.debug("Reading SQL table", table=str(full_path), format=format)

            if not hasattr(connection, "read_table") and not hasattr(connection, "read_sql_query"):
                ctx.error(
                    "Connection does not support SQL operations",
                    connection_type=type(connection).__name__,
                )
                raise ValueError(
                    f"Cannot read SQL table '{full_path}': connection type '{type(connection).__name__}' "
                    "does not support SQL operations. Use a SQL-compatible connection "
                    "(e.g., AzureSQL, PostgreSQLConnection)."
                )

            table_name = str(full_path)
            default_schema = getattr(connection, "default_schema", "dbo")
            if "." in table_name:
                schema, tbl = table_name.split(".", 1)
            else:
                schema, tbl = default_schema, table_name

            if post_read_query and hasattr(connection, "read_sql_query"):
                full_query = connection.build_select_query(tbl, schema, where=post_read_query)
                ctx.debug(
                    "Executing SQL query with incremental filter",
                    query=full_query,
                )
                return connection.read_sql_query(full_query)

            ctx.debug("Executing SQL read", schema=schema, table=tbl)
            return connection.read_table(table_name=tbl, schema=schema)
        elif format == "api":
            endpoint = str(full_path)
            ctx.debug("Reading from API", endpoint=endpoint)

            from odibi.connections.api_fetcher import create_api_fetcher
            from odibi.connections.http import HttpConnection

            if not isinstance(connection, HttpConnection):
                ctx.error(
                    "API format requires HttpConnection",
                    connection_type=type(connection).__name__,
                )
                raise ValueError(
                    f"Cannot read API data: connection type '{type(connection).__name__}' "
                    "is not an HttpConnection. Use an HTTP connection for API format."
                )

            # Extract API-specific options
            api_options = options.copy()
            params = api_options.pop("params", {})
            max_records = api_options.pop("max_records", None)
            method = api_options.pop("method", "GET")
            request_body = api_options.pop("request_body", None)

            # Create fetcher with connection's base_url and headers
            fetcher = create_api_fetcher(
                base_url=connection.base_url,
                headers=connection.headers,
                options=api_options,
            )

            # Fetch data
            df = fetcher.fetch_dataframe(
                endpoint=endpoint,
                params=params,
                max_records=max_records,
                method=method,
                request_body=request_body,
            )

            if hasattr(df, "attrs"):
                df.attrs["odibi_source_api"] = f"{connection.base_url}/{endpoint}"

            return self._process_df(df, post_read_query)
        else:
            ctx.error("Unsupported format", format=format)
            raise ValueError(
                f"Unsupported format for Pandas engine: '{format}'. "
                "Supported formats: csv, parquet, json, excel, delta, sql, api."
            )

    def write(
        self,
        df: Union[pd.DataFrame, Iterator[pd.DataFrame]],
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        register_table: Optional[str] = None,
        mode: str = "overwrite",
        options: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Write data using Pandas.

        Args:
            df: DataFrame or iterator of DataFrames to write.
            connection: Connection object providing base path and storage options.
            format: Output format (csv, parquet, delta, json, excel, avro, sql).
            table: Table name (mutually exclusive with path).
            path: File path (mutually exclusive with table).
            register_table: Optional table name to register in catalog (not used).
            mode: Write mode (overwrite, append, upsert, append_once).
            options: Additional write options to pass to pandas writers.
            streaming_config: Streaming configuration (not used in Pandas engine).

        Returns:
            Optional dict with write statistics for Delta writes, None otherwise.

        Raises:
            ValueError: If neither path nor table is provided.
        """
        ctx = get_logging_context().with_context(engine="pandas")
        start = time.time()

        destination = path or table
        ctx.debug(
            "Starting write operation",
            format=format,
            destination=destination,
            mode=mode,
        )

        # Ensure materialization if LazyDataset
        df = self.materialize(df)

        options = options or {}

        # Handle iterator/generator input
        from collections.abc import Iterator

        if isinstance(df, Iterator):
            ctx.debug("Writing iterator/generator input")
            return self._write_iterator(df, connection, format, table, path, mode, options)

        row_count = len(df)
        memory_mb = df.memory_usage(deep=True).sum() / (1024 * 1024)

        ctx.log_pandas_metrics(
            memory_mb=memory_mb,
            dtypes={col: str(dtype) for col, dtype in df.dtypes.items()},
        )

        # SQL Server / Azure SQL Support
        if is_sql_format(format):
            ctx.debug("Writing to SQL", table=table, mode=mode)
            return self._write_sql(df, connection, table, mode, options)

        # Resolve full path from connection
        try:
            full_path = self._resolve_path(path or table, connection)
        except ValueError:
            if table and not connection:
                ctx.error("Connection required when specifying 'table'", table=table)
                raise ValueError("Connection is required when specifying 'table'.")
            ctx.error("Neither path nor table provided for write operation")
            raise ValueError("Either path or table must be provided")

        # Merge storage options for cloud connections
        merged_options = self._merge_storage_options(connection, options)

        # Custom Writers
        if format in self._custom_writers:
            ctx.debug(f"Using custom writer for format: {format}")
            writer_options = merged_options.copy()
            writer_options.pop("keys", None)
            self._custom_writers[format](df, full_path, mode=mode, **writer_options)
            return None

        # Ensure directory exists (local only)
        self._ensure_directory(full_path)

        # Warn about partitioning
        self._check_partitioning(merged_options)

        # Delta Lake Write
        if format == "delta":
            ctx.debug("Writing Delta table", path=str(full_path), mode=mode)
            result = self._write_delta(df, full_path, mode, merged_options)
            elapsed = (time.time() - start) * 1000
            ctx.log_file_io(
                path=str(full_path),
                format=format,
                mode=mode,
                rows=row_count,
            )
            ctx.info(
                "Write completed",
                format=format,
                rows=row_count,
                elapsed_ms=round(elapsed, 2),
            )
            return result

        # Handle Generic Upsert/Append-Once for non-Delta
        if mode in ["upsert", "append_once"]:
            ctx.debug(f"Handling {mode} mode for non-Delta format")
            df, mode = self._handle_generic_upsert(df, full_path, format, mode, merged_options)
            row_count = len(df)

        # Standard File Write
        result = self._write_file(df, full_path, format, mode, merged_options)

        elapsed = (time.time() - start) * 1000
        ctx.log_file_io(
            path=str(full_path),
            format=format,
            mode=mode,
            rows=row_count,
        )
        ctx.info(
            "Write completed",
            format=format,
            rows=row_count,
            elapsed_ms=round(elapsed, 2),
        )

        return result

    def _write_iterator(
        self,
        df_iter: Iterator[pd.DataFrame],
        connection: Any,
        format: str,
        table: Optional[str],
        path: Optional[str],
        mode: str,
        options: Dict[str, Any],
    ) -> None:
        """Handle writing of iterator/generator."""
        first_chunk = True
        for chunk in df_iter:
            # Determine mode for this chunk
            current_mode = mode if first_chunk else "append"
            current_options = options.copy()

            # Handle CSV header for chunks
            if not first_chunk and format == "csv":
                if current_options.get("header") is not False:
                    current_options["header"] = False

            self.write(
                chunk,
                connection,
                format,
                table,
                path,
                mode=current_mode,
                options=current_options,
            )
            first_chunk = False
        return None

    def _write_sql(
        self,
        df: pd.DataFrame,
        connection: Any,
        table: Optional[str],
        mode: str,
        options: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """Handle SQL writing including merge and enhanced overwrite."""
        ctx = get_logging_context().with_context(engine="pandas")

        if not hasattr(connection, "write_table"):
            raise ValueError(
                f"Connection type '{type(connection).__name__}' does not support SQL operations"
            )

        if not table:
            raise ValueError("SQL format requires 'table' config")

        # Handle MERGE mode for SQL Server
        if mode == "merge":
            if getattr(connection, "sql_dialect", "mssql") != "mssql":
                raise NotImplementedError(
                    f"SQL MERGE mode is currently only supported for SQL Server connections. "
                    f"Connection dialect '{getattr(connection, 'sql_dialect', 'unknown')}' "
                    f"does not support MERGE. Use mode='append' or mode='overwrite' instead."
                )

            merge_keys = options.get("merge_keys")
            merge_options = options.get("merge_options")

            if not merge_keys:
                raise ValueError(
                    "MERGE mode requires 'merge_keys' in options. "
                    "Specify the key columns for the MERGE ON clause."
                )

            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            writer = SqlServerMergeWriter(connection)
            ctx.debug(
                "Executing SQL Server MERGE (Pandas)",
                target=table,
                merge_keys=merge_keys,
            )

            result = writer.merge_pandas(
                df=df,
                target_table=table,
                merge_keys=merge_keys,
                options=merge_options,
            )

            ctx.info(
                "SQL Server MERGE completed (Pandas)",
                target=table,
                inserted=result.inserted,
                updated=result.updated,
                deleted=result.deleted,
            )

            return {
                "mode": "merge",
                "inserted": result.inserted,
                "updated": result.updated,
                "deleted": result.deleted,
                "total_affected": result.total_affected,
            }

        # Handle enhanced overwrite with strategies (SQL Server only)
        if mode == "overwrite" and options.get("overwrite_options"):
            if getattr(connection, "sql_dialect", "mssql") != "mssql":
                raise NotImplementedError(
                    f"Enhanced overwrite strategies are currently only supported for SQL Server. "
                    f"Use standard mode='overwrite' without overwrite_options for "
                    f"'{getattr(connection, 'sql_dialect', 'unknown')}' connections."
                )
            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            overwrite_options = options.get("overwrite_options")
            writer = SqlServerMergeWriter(connection)

            ctx.debug(
                "Executing SQL Server enhanced overwrite (Pandas)",
                target=table,
                strategy=(
                    overwrite_options.strategy.value
                    if hasattr(overwrite_options, "strategy")
                    else "truncate_insert"
                ),
            )

            result = writer.overwrite_pandas(
                df=df,
                target_table=table,
                options=overwrite_options,
            )

            ctx.info(
                "SQL Server enhanced overwrite completed (Pandas)",
                target=table,
                strategy=result.strategy,
                rows_written=result.rows_written,
            )

            return {
                "mode": "overwrite",
                "strategy": result.strategy,
                "rows_written": result.rows_written,
            }

        # Extract schema from table name if present
        default_schema = getattr(connection, "default_schema", "dbo")
        if "." in table:
            schema, table_name = table.split(".", 1)
        else:
            schema, table_name = default_schema, table

        # Map mode to if_exists
        if_exists = "replace"  # overwrite
        if mode == "append":
            if_exists = "append"
        elif mode == "fail":
            if_exists = "fail"

        chunksize = options.get("chunksize", 1000)

        connection.write_table(
            df=df,
            table_name=table_name,
            schema=schema,
            if_exists=if_exists,
            chunksize=chunksize,
        )
        return None

    def _is_local_path(self, path: str) -> bool:
        """Check whether a path refers to a local file (not a remote URI).

        Args:
            path: File path or URI to check.

        Returns:
            True if the path is local (no scheme, ``file://``, or a Windows
            drive letter like ``D:\\``).
        """
        parsed = urlparse(str(path))
        is_windows_drive = (
            len(parsed.scheme) == 1 and parsed.scheme.isalpha() if parsed.scheme else False
        )
        return not parsed.scheme or parsed.scheme == "file" or is_windows_drive

    def _ensure_directory(self, full_path: str) -> None:
        """Ensure parent directory exists for local files."""
        if self._is_local_path(full_path):
            Path(full_path).parent.mkdir(parents=True, exist_ok=True)

    def _check_partitioning(self, options: Dict[str, Any]) -> None:
        """Warn about potential partitioning issues."""
        partition_by = options.get("partition_by") or options.get("partitionBy")
        if partition_by:
            import warnings

            warnings.warn(
                "⚠️  Partitioning can cause performance issues if misused. "
                "Only partition on low-cardinality columns (< 1000 unique values) "
                "and ensure each partition has > 1000 rows.",
                UserWarning,
            )

    def _write_delta(
        self,
        df: pd.DataFrame,
        full_path: str,
        mode: str,
        merged_options: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Handle Delta Lake writing."""
        try:
            from deltalake import DeltaTable, write_deltalake
        except ImportError:
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[pandas]' or 'pip install deltalake'. "
                "See README.md for installation instructions."
            )

        import logging

        logger = logging.getLogger(__name__)

        storage_opts = merged_options.get("storage_options", {})

        # Reset index to avoid write_deltalake including it as an extra column,
        # which causes "number of fields does not match" errors on overwrite.
        df = df.reset_index(drop=True)

        # Map modes
        delta_mode = "overwrite"
        if mode == "append":
            delta_mode = "append"
        elif mode == "error" or mode == "fail":
            delta_mode = "error"
        elif mode == "ignore":
            delta_mode = "ignore"

        # Handle null-only columns: Delta Lake doesn't support Null dtype
        # Try to match existing schema on append; fall back to string
        existing_schema = {}
        if delta_mode == "append":
            try:
                dt = DeltaTable(full_path, storage_options=storage_opts or None)
                existing_schema = {f.name: f.type for f in dt.schema().fields}
            except Exception:
                pass

        for col in df.columns:
            if df[col].isna().all():
                if col in existing_schema:
                    raw_type = str(existing_schema[col])
                    # Normalize PrimitiveType("X") → "X"
                    import re

                    m = re.match(r'PrimitiveType\("(.+)"\)', raw_type)
                    arrow_type = m.group(1) if m else raw_type
                    type_map = {
                        "int32": "Int32",
                        "int64": "Int64",
                        "long": "Int64",
                        "float": "float64",
                        "double": "float64",
                        "string": "string",
                        "boolean": "boolean",
                        "bool": "boolean",
                        "timestamp[us]": "datetime64[us]",
                        "timestamp[us, tz=UTC]": "datetime64[us, UTC]",
                        "date": "datetime64[ns]",
                    }
                    pd_type = type_map.get(arrow_type, "string")
                    if pd_type == "string" and arrow_type != "string":
                        logger.warning(
                            f"Delta write: all-null column '{col}' has unmapped Arrow type "
                            f"'{arrow_type}', falling back to string. Schema mismatch may "
                            f"occur on future writes with actual data."
                        )
                    df[col] = df[col].astype(pd_type)
                else:
                    logger.warning(
                        f"Delta write: all-null column '{col}' cast to string (no existing "
                        f"schema found). If this column should be int/float/datetime, the "
                        f"Delta schema will need schema evolution on future writes."
                    )
                    df[col] = df[col].astype("string")

        # Handle upsert/append_once logic
        if mode == "upsert":
            keys = merged_options.get("keys")
            if not keys:
                raise ValueError("Upsert requires 'keys' in options")

            if isinstance(keys, str):
                keys = [keys]

            def do_upsert():
                dt = DeltaTable(full_path, storage_options=storage_opts)
                (
                    dt.merge(
                        source=df,
                        predicate=" AND ".join([f"s.{k} = t.{k}" for k in keys]),
                        source_alias="s",
                        target_alias="t",
                    )
                    .when_matched_update_all()
                    .when_not_matched_insert_all()
                    .execute()
                )

            self._retry_delta_operation(do_upsert)
        elif mode == "append_once":
            keys = merged_options.get("keys")
            if not keys:
                raise ValueError("Append_once requires 'keys' in options")

            if isinstance(keys, str):
                keys = [keys]

            def do_append_once():
                dt = DeltaTable(full_path, storage_options=storage_opts)
                (
                    dt.merge(
                        source=df,
                        predicate=" AND ".join([f"s.{k} = t.{k}" for k in keys]),
                        source_alias="s",
                        target_alias="t",
                    )
                    .when_not_matched_insert_all()
                    .execute()
                )

            self._retry_delta_operation(do_append_once)
        else:
            # Filter options supported by write_deltalake
            write_kwargs = {
                k: v
                for k, v in merged_options.items()
                if k
                in [
                    "partition_by",
                    "mode",
                    "schema_mode",
                    "name",
                    "description",
                    "configuration",
                    "writer_properties",
                ]
            }

            # Translate legacy overwrite_schema=True to schema_mode="overwrite"
            # (overwrite_schema was removed in deltalake >= 0.18)
            if merged_options.get("overwrite_schema") and "schema_mode" not in write_kwargs:
                write_kwargs["schema_mode"] = "overwrite"

            def do_write():
                write_deltalake(
                    full_path,
                    df,
                    mode=delta_mode,
                    storage_options=storage_opts,
                    **write_kwargs,
                )

            self._retry_delta_operation(do_write)

        # Return commit info
        dt = DeltaTable(full_path, storage_options=storage_opts)
        history = dt.history(limit=1)
        latest = history[0]

        return {
            "version": dt.version(),
            "timestamp": datetime.fromtimestamp(latest.get("timestamp", 0) / 1000, tz=timezone.utc),
            "operation": latest.get("operation"),
            "operation_metrics": latest.get("operationMetrics", {}),
            "read_version": latest.get("readVersion"),
        }

    def _handle_generic_upsert(
        self,
        df: pd.DataFrame,
        full_path: str,
        format: str,
        mode: str,
        options: Dict[str, Any],
    ) -> tuple[pd.DataFrame, str]:
        """Handle upsert/append_once for standard files by merging with existing data."""
        if "keys" not in options:
            raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

        keys = options["keys"]
        if isinstance(keys, str):
            keys = [keys]

        # Try to read existing file
        existing_df = None
        try:
            read_opts = options.copy()
            read_opts.pop("keys", None)

            if format == "csv":
                existing_df = pd.read_csv(full_path, **read_opts)
            elif format == "parquet":
                existing_df = pd.read_parquet(full_path, **read_opts)
            elif format == "json":
                existing_df = pd.read_json(full_path, **read_opts)
            elif format == "excel":
                existing_df = pd.read_excel(full_path, **read_opts)
        except Exception:
            # File doesn't exist or can't be read
            return df, "overwrite"  # Treat as new write

        if existing_df is None:
            return df, "overwrite"

        if mode == "append_once":
            # Check if keys exist
            missing_keys = set(keys) - set(df.columns)
            if missing_keys:
                raise KeyError(f"Keys {missing_keys} not found in input data")

            # Identify new rows
            merged = df.merge(existing_df[keys], on=keys, how="left", indicator=True)
            new_rows = merged[merged["_merge"] == "left_only"].drop(columns=["_merge"])

            if format in ["csv", "json"]:
                return new_rows, "append"
            else:
                # Rewrite everything
                return pd.concat([existing_df, new_rows], ignore_index=True), "overwrite"

        elif mode == "upsert":
            # Check if keys exist
            missing_keys = set(keys) - set(df.columns)
            if missing_keys:
                raise KeyError(f"Keys {missing_keys} not found in input data")

            # 1. Remove rows from existing that are in input
            merged_indicator = existing_df.merge(df[keys], on=keys, how="left", indicator=True)
            rows_to_keep = existing_df[merged_indicator["_merge"] == "left_only"]

            # 2. Concat rows_to_keep + input df
            # 3. Write mode becomes overwrite
            return pd.concat([rows_to_keep, df], ignore_index=True), "overwrite"

        return df, mode

    def _write_file(
        self,
        df: pd.DataFrame,
        full_path: str,
        format: str,
        mode: str,
        merged_options: Dict[str, Any],
    ) -> None:
        """Handle standard file writing (CSV, Parquet, etc.).

        For local overwrite operations, writes to a temporary file first and
        then atomically replaces the target via ``os.replace()``.  This
        prevents data loss if the process crashes mid-write.
        """
        use_atomic = mode != "append" and self._is_local_path(full_path)

        if use_atomic:
            dir_name = os.path.dirname(full_path) or "."
            fd, tmp_path = tempfile.mkstemp(dir=dir_name, suffix=f".{format}.tmp")
            os.close(fd)
        else:
            tmp_path = None

        target = tmp_path if use_atomic else full_path

        try:
            self._write_to_target(df, target, full_path, format, mode, merged_options)
            if use_atomic:
                os.replace(tmp_path, full_path)
        except Exception:
            if tmp_path and os.path.exists(tmp_path):
                os.remove(tmp_path)
            raise

    def _write_to_target(
        self,
        df: pd.DataFrame,
        target: str,
        full_path: str,
        format: str,
        mode: str,
        merged_options: Dict[str, Any],
    ) -> None:
        """Write a DataFrame to the given target path.

        Args:
            df: DataFrame to write.
            target: Path to write to (may be a temp file for atomic writes).
            full_path: The final destination path (used for append existence
                checks).
            format: File format (csv, parquet, json, excel, avro).
            mode: Write mode (overwrite, append, etc.).
            merged_options: Writer options dict.
        """
        writer_options = merged_options.copy()
        writer_options.pop("keys", None)

        if format == "csv":
            mode_param = "w"
            if mode == "append":
                mode_param = "a"
                if not os.path.exists(full_path):
                    writer_options["header"] = True
                else:
                    if "header" not in writer_options:
                        writer_options["header"] = False

            df.to_csv(target, index=False, mode=mode_param, **writer_options)

        elif format == "parquet":
            if mode == "append" and os.path.exists(full_path):
                import pyarrow as pa
                import pyarrow.parquet as pq

                new_table = pa.Table.from_pandas(df, preserve_index=False)
                dir_name = os.path.dirname(full_path) or "."
                fd, tmp = tempfile.mkstemp(dir=dir_name, suffix=".parquet.tmp")
                os.close(fd)
                try:
                    existing_file = pq.ParquetFile(full_path)
                    existing_schema = existing_file.schema_arrow
                    unified_schema = pa.unify_schemas(
                        [existing_schema, new_table.schema], promote_options="permissive"
                    )

                    def _align_table(tbl, target_schema):
                        for field in target_schema:
                            if field.name not in tbl.schema.names:
                                tbl = tbl.append_column(field, pa.nulls(len(tbl), type=field.type))
                        return tbl.select([f.name for f in target_schema])

                    with pq.ParquetWriter(tmp, unified_schema) as writer:
                        for batch in existing_file.iter_batches():
                            table = pa.Table.from_batches([batch], schema=existing_schema)
                            writer.write_table(_align_table(table, unified_schema))
                        existing_file.close()
                        writer.write_table(_align_table(new_table, unified_schema))
                    os.replace(tmp, full_path)
                except Exception:
                    if os.path.exists(tmp):
                        try:
                            os.remove(tmp)
                        except OSError:
                            pass
                    raise
            else:
                df.to_parquet(target, index=False, **writer_options)

        elif format == "json":
            if mode == "append":
                writer_options["mode"] = "a"

            if "orient" not in writer_options:
                writer_options["orient"] = "records"

            if "storage_options" in merged_options:
                writer_options["storage_options"] = merged_options["storage_options"]

            df.to_json(target, **writer_options)

        elif format == "excel":
            if mode == "append":
                if os.path.exists(full_path):
                    with pd.ExcelWriter(full_path, mode="a", if_sheet_exists="overlay") as writer:
                        df.to_excel(writer, index=False, **writer_options)
                    return

            df.to_excel(target, index=False, **writer_options)

        elif format == "avro":
            try:
                import fastavro
            except ImportError:
                raise ImportError("Avro support requires 'pip install fastavro'")

            df_avro = df.copy()
            for col in df_avro.columns:
                if pd.api.types.is_datetime64_any_dtype(df_avro[col].dtype):
                    df_avro[col] = df_avro[col].apply(
                        lambda x: int(x.timestamp() * 1_000_000) if pd.notna(x) else None
                    )

            records = df_avro.to_dict("records")
            schema = self._infer_avro_schema(df)

            parsed = urlparse(full_path)
            if parsed.scheme and parsed.scheme not in ["file", ""]:
                import fsspec

                storage_opts = merged_options.get("storage_options", {})
                write_mode = "wb" if mode == "overwrite" else "ab"
                with fsspec.open(full_path, write_mode, **storage_opts) as f:
                    fastavro.writer(f, schema, records)
            else:
                open_mode = "wb"
                if mode == "append" and os.path.exists(full_path):
                    open_mode = "a+b"

                with open(target, open_mode) as f:
                    fastavro.writer(f, schema, records)
        else:
            raise ValueError(f"Unsupported format for Pandas engine: {format}")

    def add_write_metadata(
        self,
        df: pd.DataFrame,
        metadata_config: Any,
        source_connection: Optional[str] = None,
        source_table: Optional[str] = None,
        source_path: Optional[str] = None,
        is_file_source: bool = False,
    ) -> pd.DataFrame:
        """Add metadata columns to DataFrame before writing (Bronze layer lineage).

        Args:
            df: Pandas DataFrame
            metadata_config: WriteMetadataConfig or True (for all defaults)
            source_connection: Name of the source connection
            source_table: Name of the source table (SQL sources)
            source_path: Path of the source file (file sources)
            is_file_source: True if source is a file-based read

        Returns:
            DataFrame with metadata columns added
        """
        from odibi.config import WriteMetadataConfig

        # Normalize config: True -> all defaults
        if metadata_config is True:
            config = WriteMetadataConfig()
        elif isinstance(metadata_config, WriteMetadataConfig):
            config = metadata_config
        else:
            return df  # None or invalid -> no metadata

        # Work on a copy to avoid modifying original
        df = df.copy()

        # _extracted_at: always applicable
        if config.extracted_at:
            df["_extracted_at"] = pd.Timestamp.now()

        # _source_file: only for file sources
        if config.source_file and is_file_source and source_path:
            df["_source_file"] = source_path

        # _source_connection: all sources
        if config.source_connection and source_connection:
            df["_source_connection"] = source_connection

        # _source_table: SQL sources only
        if config.source_table and source_table:
            df["_source_table"] = source_table

        return df

    def _register_lazy_view_unused(self, conn, name: str, df: Any) -> None:
        """Register a LazyDataset as a DuckDB view."""
        duck_fmt = df.format
        if duck_fmt == "json":
            duck_fmt = "json_auto"

        if isinstance(df.path, list):
            paths = ", ".join([f"'{p}'" for p in df.path])
            conn.execute(
                f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM read_{duck_fmt}([{paths}])"
            )
        else:
            conn.execute(
                f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM read_{duck_fmt}('{df.path}')"
            )

    def execute_sql(self, sql: str, context: Context) -> pd.DataFrame:
        """Execute SQL query using DuckDB (if available) or pandasql.

        Args:
            sql: SQL query string
            context: Execution context

        Returns:
            Result DataFrame
        """
        if not isinstance(context, PandasContext):
            raise TypeError("PandasEngine requires PandasContext")

        # Try to use DuckDB for SQL
        try:
            import duckdb

            # Create in-memory database
            conn = duckdb.connect(":memory:")

            # Register all DataFrames from context
            for name in context.list_names():
                dataset_obj = context.get(name)

                # Handle LazyDataset (DuckDB optimization)
                # if isinstance(dataset_obj, LazyDataset):
                #     self._register_lazy_view(conn, name, dataset_obj)
                #     # Log that we used DuckDB on file
                #     # logger.info(f"Executing SQL via DuckDB on lazy file: {dataset_obj.path}")
                #     continue

                # Handle chunked data (Iterator)
                from collections.abc import Iterator

                if isinstance(dataset_obj, Iterator):
                    # Warning: Materializing iterator for SQL execution
                    # Note: DuckDB doesn't support streaming from iterator yet
                    chunks = [chunk for chunk in dataset_obj if not chunk.empty]
                    dataset_obj = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

                conn.register(name, dataset_obj)

            # Execute query
            result = conn.execute(sql).df()
            conn.close()

            return result

        except ImportError:
            # Fallback: try pandasql
            try:
                from pandasql import sqldf

                # Build local namespace with DataFrames
                locals_dict = {}
                for name in context.list_names():
                    df = context.get(name)

                    # Handle chunked data (Iterator)
                    from collections.abc import Iterator

                    if isinstance(df, Iterator):
                        chunks = [chunk for chunk in df if not chunk.empty]
                        df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

                    locals_dict[name] = df

                return sqldf(sql, locals_dict)

            except ImportError:
                raise TransformError(
                    "SQL execution requires 'duckdb' or 'pandasql'. "
                    "Install with: pip install duckdb"
                )

    def execute_operation(
        self,
        operation: str,
        params: Dict[str, Any],
        df: Union[pd.DataFrame, Iterator[pd.DataFrame]],
    ) -> pd.DataFrame:
        """Execute built-in operation.

        Args:
            operation: Operation name
            params: Operation parameters
            df: Input DataFrame or Iterator

        Returns:
            Result DataFrame
        """
        # Materialize LazyDataset
        df = self.materialize(df)

        # Handle chunked data (Iterator)
        from collections.abc import Iterator

        if isinstance(df, Iterator):
            # Warning: Materializing iterator for operation execution
            chunks = [chunk for chunk in df if not chunk.empty]
            df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

        if operation == "pivot":
            return self._pivot(df, params)
        elif operation == "drop_duplicates":
            return df.drop_duplicates(**params)
        elif operation == "fillna":
            return df.fillna(**params)
        elif operation == "drop":
            return df.drop(**params)
        elif operation == "rename":
            return df.rename(**params)
        elif operation == "sort":
            return df.sort_values(**params)
        elif operation == "sample":
            return df.sample(**params)
        else:
            # Fallback: check if operation is a registered transformer
            from odibi.context import EngineContext, PandasContext
            from odibi.registry import FunctionRegistry

            if FunctionRegistry.has_function(operation):
                func = FunctionRegistry.get_function(operation)
                param_model = FunctionRegistry.get_param_model(operation)

                # Create EngineContext from current df
                engine_ctx = EngineContext(
                    context=PandasContext(),
                    df=df,
                    engine=self,
                    engine_type=self.engine_type,
                    sql_executor=self.execute_sql,
                )

                # Validate and instantiate params
                if param_model:
                    validated_params = param_model(**params)
                    result_ctx = func(engine_ctx, validated_params)
                else:
                    result_ctx = func(engine_ctx, **params)

                return result_ctx.df

            raise ValueError(f"Unsupported operation: {operation}")

    def _pivot(self, df: pd.DataFrame, params: Dict[str, Any]) -> pd.DataFrame:
        """Execute pivot operation.

        Args:
            df: Input DataFrame
            params: Pivot parameters

        Returns:
            Pivoted DataFrame
        """
        group_by = params.get("group_by", [])
        pivot_column = params["pivot_column"]
        value_column = params["value_column"]
        agg_func = params.get("agg_func", "first")

        # Validate columns exist
        required_columns = set()
        if isinstance(group_by, list):
            required_columns.update(group_by)
        elif isinstance(group_by, str):
            required_columns.add(group_by)
            group_by = [group_by]

        required_columns.add(pivot_column)
        required_columns.add(value_column)

        missing = required_columns - set(df.columns)
        if missing:
            raise KeyError(
                f"Columns not found in DataFrame for pivot operation: {missing}. "
                f"Available: {list(df.columns)}"
            )

        result = df.pivot_table(
            index=group_by, columns=pivot_column, values=value_column, aggfunc=agg_func
        ).reset_index()

        # Flatten column names if multi-level
        if isinstance(result.columns, pd.MultiIndex):
            result.columns = ["_".join(col).strip("_") for col in result.columns.values]

        return result

    def harmonize_schema(
        self, df: pd.DataFrame, target_schema: Dict[str, str], policy: Any
    ) -> pd.DataFrame:
        """Harmonize DataFrame schema with target schema according to policy.

        Args:
            df: Input DataFrame to harmonize.
            target_schema: Target schema as dict of column_name -> type_string.
            policy: Schema policy defining how to handle missing/new columns.

        Returns:
            DataFrame with schema harmonized to match target_schema.

        Raises:
            ValueError: If policy validation fails for missing or new columns.
        """
        # Ensure materialization
        df = self.materialize(df)

        from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

        target_cols = list(target_schema.keys())
        current_cols = df.columns.tolist()

        missing = set(target_cols) - set(current_cols)
        new_cols = set(current_cols) - set(target_cols)

        # 1. Check Validations
        if missing and policy.on_missing_columns == OnMissingColumns.FAIL:
            raise ValueError(f"Schema Policy Violation: Missing columns {missing}")

        if new_cols and policy.on_new_columns == OnNewColumns.FAIL:
            raise ValueError(f"Schema Policy Violation: New columns {new_cols}")

        # 2. Apply Transformations
        if policy.mode == SchemaMode.EVOLVE and policy.on_new_columns == OnNewColumns.ADD_NULLABLE:
            # Evolve: Add missing columns, Keep new columns
            for col in missing:
                df[col] = None
        else:
            # Enforce / Ignore New: Project to target schema (Drops new, Adds missing)
            # Note: reindex adds NaN for missing columns
            df = df.reindex(columns=target_cols)

        return df

    def anonymize(
        self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
    ) -> pd.DataFrame:
        """Anonymize specified columns.

        Args:
            df: DataFrame or LazyDataset to anonymize.
            columns: List of column names to anonymize.
            method: Anonymization method (hash, mask, or redact).
            salt: Optional salt string for hash method.

        Returns:
            DataFrame with specified columns anonymized.
        """
        # Ensure materialization
        df = self.materialize(df)

        res = df.copy()

        for col in columns:
            if col not in res.columns:
                continue

            if method == "hash":
                # Vectorized Hashing (via map/apply)
                # Note: True vectorization requires C-level support (e.g. pyarrow.compute)
                # Standard Pandas apply is the fallback but we can optimize string handling

                # Convert to string, handling nulls
                # s_col = res[col].astype(str)
                # Nulls become 'nan'/'None' string, we want to preserve them or hash them consistently?
                # Typically nulls should remain null.

                mask_nulls = res[col].isna()

                def _hash_val(val):
                    to_hash = val
                    if salt:
                        to_hash += salt
                    return hashlib.sha256(to_hash.encode("utf-8")).hexdigest()

                # Apply only to non-nulls
                res.loc[~mask_nulls, col] = res.loc[~mask_nulls, col].astype(str).apply(_hash_val)

            elif method == "mask":
                # Vectorized Masking
                # Mask all but last 4 characters

                mask_nulls = res[col].isna()
                s_valid = res.loc[~mask_nulls, col].astype(str)

                # Use vectorized regex replacement
                # Replace any character that is followed by 4 characters with '*'
                res.loc[~mask_nulls, col] = s_valid.str.replace(r".(?=.{4})", "*", regex=True)

            elif method == "redact":
                res[col] = "[REDACTED]"

        return res

    def get_schema(self, df: Any) -> Dict[str, str]:
        """Get DataFrame schema with types.

        Args:
            df: DataFrame or LazyDataset

        Returns:
            Dict[str, str]: Column name -> Type string
        """
        if isinstance(df, LazyDataset):
            if self.use_duckdb:
                try:
                    import duckdb

                    conn = duckdb.connect(":memory:")
                    self._register_lazy_view(conn, "df", df)
                    res = conn.execute("DESCRIBE SELECT * FROM df").df()
                    return dict(zip(res["column_name"], res["column_type"]))
                except Exception:
                    pass
            df = self.materialize(df)

        return {col: str(df[col].dtype) for col in df.columns}

    def get_shape(self, df: Any) -> tuple:
        """Get DataFrame shape.

        Args:
            df: DataFrame or LazyDataset

        Returns:
            (rows, columns)
        """
        if isinstance(df, LazyDataset):
            cols = len(self.get_schema(df))
            rows = self.count_rows(df)
            return (rows, cols)
        return df.shape

    def count_rows(self, df: Any) -> int:
        """Count rows in DataFrame.

        Args:
            df: DataFrame or LazyDataset

        Returns:
            Row count
        """
        if isinstance(df, LazyDataset):
            if self.use_duckdb:
                try:
                    import duckdb

                    conn = duckdb.connect(":memory:")
                    self._register_lazy_view(conn, "df", df)
                    res = conn.execute("SELECT count(*) FROM df").fetchone()
                    return res[0] if res else 0
                except Exception:
                    pass
            df = self.materialize(df)

        return len(df)

    def count_nulls(self, df: pd.DataFrame, columns: List[str]) -> Dict[str, int]:
        """Count nulls in specified columns.

        Args:
            df: DataFrame
            columns: Columns to check

        Returns:
            Dictionary of column -> null count
        """
        null_counts = {}
        for col in columns:
            if col in df.columns:
                null_counts[col] = int(df[col].isna().sum())
            else:
                raise ValueError(
                    f"Column '{col}' not found in DataFrame. Available columns: {list(df.columns)}"
                )
        return null_counts

    def validate_schema(self, df: pd.DataFrame, schema_rules: Dict[str, Any]) -> List[str]:
        """Validate DataFrame schema.

        Args:
            df: DataFrame
            schema_rules: Validation rules

        Returns:
            List of validation failures
        """
        # Ensure materialization
        df = self.materialize(df)

        failures = []

        # Check required columns
        if "required_columns" in schema_rules:
            required = schema_rules["required_columns"]
            missing = set(required) - set(df.columns)
            if missing:
                failures.append(f"Missing required columns: {', '.join(missing)}")

        # Check column types
        if "types" in schema_rules:
            type_map = {
                "int": ["int64", "int32", "int16", "int8"],
                "float": ["float64", "float32"],
                "str": ["object", "string"],
                "bool": ["bool"],
            }

            for col, expected_type in schema_rules["types"].items():
                if col not in df.columns:
                    failures.append(f"Column '{col}' not found for type validation")
                    continue

                actual_type = str(df[col].dtype)
                # Handle pyarrow types (e.g. int64[pyarrow])
                if "[" in actual_type and "pyarrow" in actual_type:
                    actual_type = actual_type.split("[")[0]

                expected_dtypes = type_map.get(expected_type, [expected_type])

                if actual_type not in expected_dtypes:
                    failures.append(
                        f"Column '{col}' has type '{actual_type}', expected '{expected_type}'"
                    )

        return failures

    def _infer_avro_schema(self, df: pd.DataFrame) -> Dict[str, Any]:
        """Infer Avro schema from pandas DataFrame.

        Args:
            df: DataFrame to infer schema from

        Returns:
            Avro schema dictionary
        """
        type_mapping = {
            "int64": "long",
            "int32": "int",
            "float64": "double",
            "float32": "float",
            "bool": "boolean",
            "object": "string",
            "string": "string",
        }

        fields = []
        for col in df.columns:
            dtype = df[col].dtype
            dtype_str = str(dtype)

            # Handle datetime types with Avro logical types
            if pd.api.types.is_datetime64_any_dtype(dtype):
                avro_type = {
                    "type": "long",
                    "logicalType": "timestamp-micros",
                }
            elif dtype_str == "date" or (hasattr(dtype, "name") and "date" in dtype.name.lower()):
                avro_type = {
                    "type": "int",
                    "logicalType": "date",
                }
            elif pd.api.types.is_timedelta64_dtype(dtype):
                avro_type = {
                    "type": "long",
                    "logicalType": "time-micros",
                }
            else:
                avro_type = type_mapping.get(dtype_str, "string")

            # Handle nullable columns
            if df[col].isnull().any():
                avro_type = ["null", avro_type]

            fields.append({"name": col, "type": avro_type})

        return {"type": "record", "name": "DataFrame", "fields": fields}

    def validate_data(self, df: pd.DataFrame, validation_config: Any) -> List[str]:
        """Validate DataFrame against rules.

        Args:
            df: DataFrame
            validation_config: ValidationConfig object

        Returns:
            List of validation failure messages
        """
        # Ensure materialization
        df = self.materialize(df)

        failures = []

        # Check not empty
        if validation_config.not_empty:
            if len(df) == 0:
                failures.append("DataFrame is empty")

        # Check for nulls in specified columns
        if validation_config.no_nulls:
            null_counts = self.count_nulls(df, validation_config.no_nulls)
            for col, count in null_counts.items():
                if count > 0:
                    failures.append(f"Column '{col}' has {count} null values")

        # Schema validation
        if validation_config.schema_validation:
            schema_failures = self.validate_schema(df, validation_config.schema_validation)
            failures.extend(schema_failures)

        # Range validation
        if validation_config.ranges:
            for col, bounds in validation_config.ranges.items():
                if col in df.columns:
                    min_val = bounds.get("min")
                    max_val = bounds.get("max")

                    if min_val is not None:
                        min_violations = df[df[col] < min_val]
                        if len(min_violations) > 0:
                            failures.append(f"Column '{col}' has values < {min_val}")

                    if max_val is not None:
                        max_violations = df[df[col] > max_val]
                        if len(max_violations) > 0:
                            failures.append(f"Column '{col}' has values > {max_val}")
                else:
                    failures.append(f"Column '{col}' not found for range validation")

        # Allowed values validation
        if validation_config.allowed_values:
            for col, allowed in validation_config.allowed_values.items():
                if col in df.columns:
                    # Check for values not in allowed list
                    invalid = df[~df[col].isin(allowed)]
                    if len(invalid) > 0:
                        failures.append(f"Column '{col}' has invalid values")
                else:
                    failures.append(f"Column '{col}' not found for allowed values validation")

        return failures

    def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
        """Get sample rows as list of dictionaries.

        Args:
            df: DataFrame or LazyDataset
            n: Number of rows to return

        Returns:
            List of row dictionaries
        """
        if isinstance(df, LazyDataset):
            if self.use_duckdb:
                try:
                    import duckdb

                    conn = duckdb.connect(":memory:")
                    self._register_lazy_view(conn, "df", df)
                    res_df = conn.execute(f"SELECT * FROM df LIMIT {n}").df()
                    return res_df.to_dict("records")
                except Exception:
                    pass
            df = self.materialize(df)

        return df.head(n).to_dict("records")

    def table_exists(
        self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
    ) -> bool:
        """Check if table or location exists.

        Args:
            connection: Connection object
            table: Table name (not used in Pandas—no catalog)
            path: File path

        Returns:
            True if file/directory exists, False otherwise
        """
        if path:
            full_path = connection.get_path(path)
            return os.path.exists(full_path)
        return False

    def get_table_schema(
        self,
        connection: Any,
        table: Optional[str] = None,
        path: Optional[str] = None,
        format: Optional[str] = None,
    ) -> Optional[Dict[str, str]]:
        """Get schema of an existing table/file.

        Args:
            connection: Connection object providing base path and storage options.
            table: Table name (for SQL sources).
            path: File path (for file sources).
            format: Data format (delta, parquet, csv, sql, sql_server, azure_sql).

        Returns:
            Dict mapping column names to type strings, or None if schema
            cannot be inferred.
        """
        try:
            if table and is_sql_format(format):
                query = connection.build_select_query(table, limit=0)
                df = connection.read_sql(query)
                return self.get_schema(df)

            if path:
                full_path = connection.get_path(path)
                if not os.path.exists(full_path):
                    return None

                if format == "delta":
                    from deltalake import DeltaTable

                    dt = DeltaTable(full_path)
                    # Use pyarrow schema to pandas schema to avoid reading data
                    arrow_schema = dt.schema().to_pyarrow()
                    empty_df = arrow_schema.empty_table().to_pandas()
                    return self.get_schema(empty_df)

                elif format == "parquet":
                    import pyarrow.parquet as pq

                    target_path = full_path
                    if os.path.isdir(full_path):
                        # Find first parquet file
                        files = glob.glob(os.path.join(full_path, "*.parquet"))
                        if not files:
                            return None
                        target_path = files[0]

                    schema = pq.read_schema(target_path)
                    empty_df = schema.empty_table().to_pandas()
                    return self.get_schema(empty_df)

                elif format == "csv":
                    df = pd.read_csv(full_path, nrows=0)
                    return self.get_schema(df)

        except (FileNotFoundError, PermissionError):
            return None
        except ImportError as e:
            # Log missing optional dependency
            import logging

            logging.getLogger(__name__).warning(
                f"Could not infer schema due to missing dependency: {e}"
            )
            return None
        except Exception as e:
            import logging

            logging.getLogger(__name__).warning(f"Failed to infer schema for {table or path}: {e}")
            return None
        return None

    def vacuum_delta(
        self,
        connection: Any,
        path: str,
        retention_hours: int = 168,
        dry_run: bool = False,
        enforce_retention_duration: bool = True,
    ) -> Dict[str, Any]:
        """VACUUM a Delta table to remove old files.

        Args:
            connection: Connection object
            path: Delta table path
            retention_hours: Retention period (default 168 = 7 days)
            dry_run: If True, only show files to be deleted
            enforce_retention_duration: If False, allows retention < 168 hours (testing only)

        Returns:
            Dictionary with files_deleted count
        """
        ctx = get_logging_context().with_context(engine="pandas")
        start = time.time()

        ctx.debug(
            "Starting Delta VACUUM",
            path=path,
            retention_hours=retention_hours,
            dry_run=dry_run,
        )

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[pandas]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        deleted_files = dt.vacuum(
            retention_hours=retention_hours,
            dry_run=dry_run,
            enforce_retention_duration=enforce_retention_duration,
        )

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta VACUUM completed",
            path=str(full_path),
            files_deleted=len(deleted_files),
            dry_run=dry_run,
            elapsed_ms=round(elapsed, 2),
        )

        return {"files_deleted": len(deleted_files)}

    def get_delta_history(
        self, connection: Any, path: str, limit: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Get Delta table history.

        Args:
            connection: Connection object
            path: Delta table path
            limit: Maximum number of versions to return

        Returns:
            List of version metadata dictionaries
        """
        ctx = get_logging_context().with_context(engine="pandas")
        start = time.time()

        ctx.debug("Getting Delta table history", path=path, limit=limit)

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[pandas]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        history = dt.history(limit=limit)

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta history retrieved",
            path=str(full_path),
            versions_returned=len(history) if history else 0,
            elapsed_ms=round(elapsed, 2),
        )

        return history

    def restore_delta(self, connection: Any, path: str, version: int) -> None:
        """Restore Delta table to a specific version.

        Args:
            connection: Connection object
            path: Delta table path
            version: Version number to restore to
        """
        ctx = get_logging_context().with_context(engine="pandas")
        start = time.time()

        ctx.info("Starting Delta table restore", path=path, target_version=version)

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[pandas]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        dt.restore(version)

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta table restored",
            path=str(full_path),
            restored_to_version=version,
            elapsed_ms=round(elapsed, 2),
        )

    def maintain_table(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        config: Optional[Any] = None,
    ) -> None:
        """Run table maintenance operations (optimize, vacuum).

        Args:
            connection: Connection object providing base path and storage options.
            format: Data format (only delta is supported).
            table: Table name (mutually exclusive with path).
            path: File path (mutually exclusive with table).
            config: Maintenance configuration specifying operations to run.
        """
        ctx = get_logging_context().with_context(engine="pandas")

        if format != "delta" or not config or not config.enabled:
            return

        if not path and not table:
            return

        full_path = connection.get_path(path if path else table)
        start = time.time()

        ctx.info("Starting table maintenance", path=str(full_path))

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.warning(
                "Auto-optimize skipped: 'deltalake' library not installed",
                path=str(full_path),
            )
            return

        try:
            storage_opts = {}
            if hasattr(connection, "pandas_storage_options"):
                storage_opts = connection.pandas_storage_options()

            dt = DeltaTable(full_path, storage_options=storage_opts)

            ctx.info("Running Delta OPTIMIZE (compaction)", path=str(full_path))
            dt.optimize.compact()

            retention = config.vacuum_retention_hours
            if retention is not None and retention > 0:
                ctx.info(
                    "Running Delta VACUUM",
                    path=str(full_path),
                    retention_hours=retention,
                )
                dt.vacuum(
                    retention_hours=retention,
                    enforce_retention_duration=True,
                    dry_run=False,
                )

            elapsed = (time.time() - start) * 1000
            ctx.info(
                "Table maintenance completed",
                path=str(full_path),
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            ctx.warning(
                "Auto-optimize failed",
                path=str(full_path),
                error=str(e),
            )

    def get_source_files(self, df: Any) -> List[str]:
        """Get list of source files that generated this DataFrame.

        Args:
            df: DataFrame or LazyDataset

        Returns:
            List of file paths
        """
        if isinstance(df, LazyDataset):
            if isinstance(df.path, list):
                return df.path
            return [str(df.path)]

        if hasattr(df, "attrs"):
            return df.attrs.get("odibi_source_files", [])
        return []

    def profile_nulls(self, df: pd.DataFrame) -> Dict[str, float]:
        """Calculate null percentage for each column.

        Args:
            df: DataFrame

        Returns:
            Dictionary of {column_name: null_percentage}
        """
        # Ensure materialization
        df = self.materialize(df)

        # mean() of boolean DataFrame gives the percentage of True values
        return df.isna().mean().to_dict()

    def filter_greater_than(self, df: pd.DataFrame, column: str, value: Any) -> pd.DataFrame:
        """Filter DataFrame where column > value.

        Automatically casts string columns to datetime for proper comparison.
        """
        if column not in df.columns:
            raise ValueError(f"Column '{column}' not found in DataFrame")

        try:
            col_series = df[column]

            if pd.api.types.is_string_dtype(col_series):
                col_series = pd.to_datetime(col_series, errors="coerce")
            elif pd.api.types.is_datetime64_any_dtype(col_series) and isinstance(value, str):
                value = pd.to_datetime(value)

            return df[col_series > value]
        except Exception as e:
            raise ValueError(f"Failed to filter {column} > {value}: {e}")

    def filter_coalesce(
        self, df: pd.DataFrame, col1: str, col2: str, op: str, value: Any
    ) -> pd.DataFrame:
        """Filter using COALESCE(col1, col2) op value.

        Automatically casts string columns to datetime for proper comparison.
        """
        if col1 not in df.columns:
            raise ValueError(f"Column '{col1}' not found")

        def _to_datetime_if_string(series: pd.Series) -> pd.Series:
            if pd.api.types.is_string_dtype(series):
                return pd.to_datetime(series, errors="coerce")
            return series

        s1 = _to_datetime_if_string(df[col1])

        if col2 not in df.columns:
            s = s1
        else:
            s2 = _to_datetime_if_string(df[col2])
            s = s1.combine_first(s2)

        try:
            if pd.api.types.is_datetime64_any_dtype(s) and isinstance(value, str):
                value = pd.to_datetime(value)

            if op == ">=":
                return df[s >= value]
            elif op == ">":
                return df[s > value]
            elif op == "<=":
                return df[s <= value]
            elif op == "<":
                return df[s < value]
            elif op == "==" or op == "=":
                return df[s == value]
            else:
                raise ValueError(f"Unsupported operator: {op}")
        except Exception as e:
            raise ValueError(f"Failed to filter COALESCE({col1}, {col2}) {op} {value}: {e}")

__init__(connections=None, config=None)

Initialize Pandas engine.

Parameters:

Name Type Description Default
connections Optional[Dict[str, Any]]

Dictionary of connection objects

None
config Optional[Dict[str, Any]]

Engine configuration (optional)

None
Source code in odibi/engine/pandas_engine.py
def __init__(
    self,
    connections: Optional[Dict[str, Any]] = None,
    config: Optional[Dict[str, Any]] = None,
):
    """Initialize Pandas engine.

    Args:
        connections: Dictionary of connection objects
        config: Engine configuration (optional)
    """
    self.connections = connections or {}
    self.config = config or {}

    # Suppress noisy delta-rs transaction conflict warnings (handled by retry)
    if "RUST_LOG" not in os.environ:
        os.environ["RUST_LOG"] = "deltalake_core::kernel::transaction=error"

    # Check for performance flags
    performance = self.config.get("performance", {})

    # Determine desired state
    if hasattr(performance, "use_arrow"):
        desired_use_arrow = performance.use_arrow
    elif isinstance(performance, dict):
        desired_use_arrow = performance.get("use_arrow", True)
    else:
        desired_use_arrow = True

    # Verify availability
    if desired_use_arrow:
        try:
            import pyarrow  # noqa: F401

            self.use_arrow = True
        except ImportError:
            import logging

            logger = logging.getLogger(__name__)
            logger.warning(
                "Apache Arrow not found. Disabling Arrow optimizations. "
                "Install 'pyarrow' to enable."
            )
            self.use_arrow = False
    else:
        self.use_arrow = False

    # Check for DuckDB
    self.use_duckdb = False
    # Default to False to ensure stability with existing tests (Lazy Loading is opt-in)
    if self.config.get("performance", {}).get("use_duckdb", False):
        try:
            import duckdb  # noqa: F401

            self.use_duckdb = True
        except ImportError:
            pass

add_write_metadata(df, metadata_config, source_connection=None, source_table=None, source_path=None, is_file_source=False)

Add metadata columns to DataFrame before writing (Bronze layer lineage).

Parameters:

Name Type Description Default
df DataFrame

Pandas DataFrame

required
metadata_config Any

WriteMetadataConfig or True (for all defaults)

required
source_connection Optional[str]

Name of the source connection

None
source_table Optional[str]

Name of the source table (SQL sources)

None
source_path Optional[str]

Path of the source file (file sources)

None
is_file_source bool

True if source is a file-based read

False

Returns:

Type Description
DataFrame

DataFrame with metadata columns added

Source code in odibi/engine/pandas_engine.py
def add_write_metadata(
    self,
    df: pd.DataFrame,
    metadata_config: Any,
    source_connection: Optional[str] = None,
    source_table: Optional[str] = None,
    source_path: Optional[str] = None,
    is_file_source: bool = False,
) -> pd.DataFrame:
    """Add metadata columns to DataFrame before writing (Bronze layer lineage).

    Args:
        df: Pandas DataFrame
        metadata_config: WriteMetadataConfig or True (for all defaults)
        source_connection: Name of the source connection
        source_table: Name of the source table (SQL sources)
        source_path: Path of the source file (file sources)
        is_file_source: True if source is a file-based read

    Returns:
        DataFrame with metadata columns added
    """
    from odibi.config import WriteMetadataConfig

    # Normalize config: True -> all defaults
    if metadata_config is True:
        config = WriteMetadataConfig()
    elif isinstance(metadata_config, WriteMetadataConfig):
        config = metadata_config
    else:
        return df  # None or invalid -> no metadata

    # Work on a copy to avoid modifying original
    df = df.copy()

    # _extracted_at: always applicable
    if config.extracted_at:
        df["_extracted_at"] = pd.Timestamp.now()

    # _source_file: only for file sources
    if config.source_file and is_file_source and source_path:
        df["_source_file"] = source_path

    # _source_connection: all sources
    if config.source_connection and source_connection:
        df["_source_connection"] = source_connection

    # _source_table: SQL sources only
    if config.source_table and source_table:
        df["_source_table"] = source_table

    return df

anonymize(df, columns, method, salt=None)

Anonymize specified columns.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset to anonymize.

required
columns List[str]

List of column names to anonymize.

required
method str

Anonymization method (hash, mask, or redact).

required
salt Optional[str]

Optional salt string for hash method.

None

Returns:

Type Description
DataFrame

DataFrame with specified columns anonymized.

Source code in odibi/engine/pandas_engine.py
def anonymize(
    self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
) -> pd.DataFrame:
    """Anonymize specified columns.

    Args:
        df: DataFrame or LazyDataset to anonymize.
        columns: List of column names to anonymize.
        method: Anonymization method (hash, mask, or redact).
        salt: Optional salt string for hash method.

    Returns:
        DataFrame with specified columns anonymized.
    """
    # Ensure materialization
    df = self.materialize(df)

    res = df.copy()

    for col in columns:
        if col not in res.columns:
            continue

        if method == "hash":
            # Vectorized Hashing (via map/apply)
            # Note: True vectorization requires C-level support (e.g. pyarrow.compute)
            # Standard Pandas apply is the fallback but we can optimize string handling

            # Convert to string, handling nulls
            # s_col = res[col].astype(str)
            # Nulls become 'nan'/'None' string, we want to preserve them or hash them consistently?
            # Typically nulls should remain null.

            mask_nulls = res[col].isna()

            def _hash_val(val):
                to_hash = val
                if salt:
                    to_hash += salt
                return hashlib.sha256(to_hash.encode("utf-8")).hexdigest()

            # Apply only to non-nulls
            res.loc[~mask_nulls, col] = res.loc[~mask_nulls, col].astype(str).apply(_hash_val)

        elif method == "mask":
            # Vectorized Masking
            # Mask all but last 4 characters

            mask_nulls = res[col].isna()
            s_valid = res.loc[~mask_nulls, col].astype(str)

            # Use vectorized regex replacement
            # Replace any character that is followed by 4 characters with '*'
            res.loc[~mask_nulls, col] = s_valid.str.replace(r".(?=.{4})", "*", regex=True)

        elif method == "redact":
            res[col] = "[REDACTED]"

    return res

count_nulls(df, columns)

Count nulls in specified columns.

Parameters:

Name Type Description Default
df DataFrame

DataFrame

required
columns List[str]

Columns to check

required

Returns:

Type Description
Dict[str, int]

Dictionary of column -> null count

Source code in odibi/engine/pandas_engine.py
def count_nulls(self, df: pd.DataFrame, columns: List[str]) -> Dict[str, int]:
    """Count nulls in specified columns.

    Args:
        df: DataFrame
        columns: Columns to check

    Returns:
        Dictionary of column -> null count
    """
    null_counts = {}
    for col in columns:
        if col in df.columns:
            null_counts[col] = int(df[col].isna().sum())
        else:
            raise ValueError(
                f"Column '{col}' not found in DataFrame. Available columns: {list(df.columns)}"
            )
    return null_counts

count_rows(df)

Count rows in DataFrame.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required

Returns:

Type Description
int

Row count

Source code in odibi/engine/pandas_engine.py
def count_rows(self, df: Any) -> int:
    """Count rows in DataFrame.

    Args:
        df: DataFrame or LazyDataset

    Returns:
        Row count
    """
    if isinstance(df, LazyDataset):
        if self.use_duckdb:
            try:
                import duckdb

                conn = duckdb.connect(":memory:")
                self._register_lazy_view(conn, "df", df)
                res = conn.execute("SELECT count(*) FROM df").fetchone()
                return res[0] if res else 0
            except Exception:
                pass
        df = self.materialize(df)

    return len(df)

execute_operation(operation, params, df)

Execute built-in operation.

Parameters:

Name Type Description Default
operation str

Operation name

required
params Dict[str, Any]

Operation parameters

required
df Union[DataFrame, Iterator[DataFrame]]

Input DataFrame or Iterator

required

Returns:

Type Description
DataFrame

Result DataFrame

Source code in odibi/engine/pandas_engine.py
def execute_operation(
    self,
    operation: str,
    params: Dict[str, Any],
    df: Union[pd.DataFrame, Iterator[pd.DataFrame]],
) -> pd.DataFrame:
    """Execute built-in operation.

    Args:
        operation: Operation name
        params: Operation parameters
        df: Input DataFrame or Iterator

    Returns:
        Result DataFrame
    """
    # Materialize LazyDataset
    df = self.materialize(df)

    # Handle chunked data (Iterator)
    from collections.abc import Iterator

    if isinstance(df, Iterator):
        # Warning: Materializing iterator for operation execution
        chunks = [chunk for chunk in df if not chunk.empty]
        df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

    if operation == "pivot":
        return self._pivot(df, params)
    elif operation == "drop_duplicates":
        return df.drop_duplicates(**params)
    elif operation == "fillna":
        return df.fillna(**params)
    elif operation == "drop":
        return df.drop(**params)
    elif operation == "rename":
        return df.rename(**params)
    elif operation == "sort":
        return df.sort_values(**params)
    elif operation == "sample":
        return df.sample(**params)
    else:
        # Fallback: check if operation is a registered transformer
        from odibi.context import EngineContext, PandasContext
        from odibi.registry import FunctionRegistry

        if FunctionRegistry.has_function(operation):
            func = FunctionRegistry.get_function(operation)
            param_model = FunctionRegistry.get_param_model(operation)

            # Create EngineContext from current df
            engine_ctx = EngineContext(
                context=PandasContext(),
                df=df,
                engine=self,
                engine_type=self.engine_type,
                sql_executor=self.execute_sql,
            )

            # Validate and instantiate params
            if param_model:
                validated_params = param_model(**params)
                result_ctx = func(engine_ctx, validated_params)
            else:
                result_ctx = func(engine_ctx, **params)

            return result_ctx.df

        raise ValueError(f"Unsupported operation: {operation}")

execute_sql(sql, context)

Execute SQL query using DuckDB (if available) or pandasql.

Parameters:

Name Type Description Default
sql str

SQL query string

required
context Context

Execution context

required

Returns:

Type Description
DataFrame

Result DataFrame

Source code in odibi/engine/pandas_engine.py
def execute_sql(self, sql: str, context: Context) -> pd.DataFrame:
    """Execute SQL query using DuckDB (if available) or pandasql.

    Args:
        sql: SQL query string
        context: Execution context

    Returns:
        Result DataFrame
    """
    if not isinstance(context, PandasContext):
        raise TypeError("PandasEngine requires PandasContext")

    # Try to use DuckDB for SQL
    try:
        import duckdb

        # Create in-memory database
        conn = duckdb.connect(":memory:")

        # Register all DataFrames from context
        for name in context.list_names():
            dataset_obj = context.get(name)

            # Handle LazyDataset (DuckDB optimization)
            # if isinstance(dataset_obj, LazyDataset):
            #     self._register_lazy_view(conn, name, dataset_obj)
            #     # Log that we used DuckDB on file
            #     # logger.info(f"Executing SQL via DuckDB on lazy file: {dataset_obj.path}")
            #     continue

            # Handle chunked data (Iterator)
            from collections.abc import Iterator

            if isinstance(dataset_obj, Iterator):
                # Warning: Materializing iterator for SQL execution
                # Note: DuckDB doesn't support streaming from iterator yet
                chunks = [chunk for chunk in dataset_obj if not chunk.empty]
                dataset_obj = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

            conn.register(name, dataset_obj)

        # Execute query
        result = conn.execute(sql).df()
        conn.close()

        return result

    except ImportError:
        # Fallback: try pandasql
        try:
            from pandasql import sqldf

            # Build local namespace with DataFrames
            locals_dict = {}
            for name in context.list_names():
                df = context.get(name)

                # Handle chunked data (Iterator)
                from collections.abc import Iterator

                if isinstance(df, Iterator):
                    chunks = [chunk for chunk in df if not chunk.empty]
                    df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()

                locals_dict[name] = df

            return sqldf(sql, locals_dict)

        except ImportError:
            raise TransformError(
                "SQL execution requires 'duckdb' or 'pandasql'. "
                "Install with: pip install duckdb"
            )

filter_coalesce(df, col1, col2, op, value)

Filter using COALESCE(col1, col2) op value.

Automatically casts string columns to datetime for proper comparison.

Source code in odibi/engine/pandas_engine.py
def filter_coalesce(
    self, df: pd.DataFrame, col1: str, col2: str, op: str, value: Any
) -> pd.DataFrame:
    """Filter using COALESCE(col1, col2) op value.

    Automatically casts string columns to datetime for proper comparison.
    """
    if col1 not in df.columns:
        raise ValueError(f"Column '{col1}' not found")

    def _to_datetime_if_string(series: pd.Series) -> pd.Series:
        if pd.api.types.is_string_dtype(series):
            return pd.to_datetime(series, errors="coerce")
        return series

    s1 = _to_datetime_if_string(df[col1])

    if col2 not in df.columns:
        s = s1
    else:
        s2 = _to_datetime_if_string(df[col2])
        s = s1.combine_first(s2)

    try:
        if pd.api.types.is_datetime64_any_dtype(s) and isinstance(value, str):
            value = pd.to_datetime(value)

        if op == ">=":
            return df[s >= value]
        elif op == ">":
            return df[s > value]
        elif op == "<=":
            return df[s <= value]
        elif op == "<":
            return df[s < value]
        elif op == "==" or op == "=":
            return df[s == value]
        else:
            raise ValueError(f"Unsupported operator: {op}")
    except Exception as e:
        raise ValueError(f"Failed to filter COALESCE({col1}, {col2}) {op} {value}: {e}")

filter_greater_than(df, column, value)

Filter DataFrame where column > value.

Automatically casts string columns to datetime for proper comparison.

Source code in odibi/engine/pandas_engine.py
def filter_greater_than(self, df: pd.DataFrame, column: str, value: Any) -> pd.DataFrame:
    """Filter DataFrame where column > value.

    Automatically casts string columns to datetime for proper comparison.
    """
    if column not in df.columns:
        raise ValueError(f"Column '{column}' not found in DataFrame")

    try:
        col_series = df[column]

        if pd.api.types.is_string_dtype(col_series):
            col_series = pd.to_datetime(col_series, errors="coerce")
        elif pd.api.types.is_datetime64_any_dtype(col_series) and isinstance(value, str):
            value = pd.to_datetime(value)

        return df[col_series > value]
    except Exception as e:
        raise ValueError(f"Failed to filter {column} > {value}: {e}")

get_delta_history(connection, path, limit=None)

Get Delta table history.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
limit Optional[int]

Maximum number of versions to return

None

Returns:

Type Description
List[Dict[str, Any]]

List of version metadata dictionaries

Source code in odibi/engine/pandas_engine.py
def get_delta_history(
    self, connection: Any, path: str, limit: Optional[int] = None
) -> List[Dict[str, Any]]:
    """Get Delta table history.

    Args:
        connection: Connection object
        path: Delta table path
        limit: Maximum number of versions to return

    Returns:
        List of version metadata dictionaries
    """
    ctx = get_logging_context().with_context(engine="pandas")
    start = time.time()

    ctx.debug("Getting Delta table history", path=path, limit=limit)

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[pandas]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    history = dt.history(limit=limit)

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta history retrieved",
        path=str(full_path),
        versions_returned=len(history) if history else 0,
        elapsed_ms=round(elapsed, 2),
    )

    return history

get_sample(df, n=10)

Get sample rows as list of dictionaries.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required
n int

Number of rows to return

10

Returns:

Type Description
List[Dict[str, Any]]

List of row dictionaries

Source code in odibi/engine/pandas_engine.py
def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
    """Get sample rows as list of dictionaries.

    Args:
        df: DataFrame or LazyDataset
        n: Number of rows to return

    Returns:
        List of row dictionaries
    """
    if isinstance(df, LazyDataset):
        if self.use_duckdb:
            try:
                import duckdb

                conn = duckdb.connect(":memory:")
                self._register_lazy_view(conn, "df", df)
                res_df = conn.execute(f"SELECT * FROM df LIMIT {n}").df()
                return res_df.to_dict("records")
            except Exception:
                pass
        df = self.materialize(df)

    return df.head(n).to_dict("records")

get_schema(df)

Get DataFrame schema with types.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required

Returns:

Type Description
Dict[str, str]

Dict[str, str]: Column name -> Type string

Source code in odibi/engine/pandas_engine.py
def get_schema(self, df: Any) -> Dict[str, str]:
    """Get DataFrame schema with types.

    Args:
        df: DataFrame or LazyDataset

    Returns:
        Dict[str, str]: Column name -> Type string
    """
    if isinstance(df, LazyDataset):
        if self.use_duckdb:
            try:
                import duckdb

                conn = duckdb.connect(":memory:")
                self._register_lazy_view(conn, "df", df)
                res = conn.execute("DESCRIBE SELECT * FROM df").df()
                return dict(zip(res["column_name"], res["column_type"]))
            except Exception:
                pass
        df = self.materialize(df)

    return {col: str(df[col].dtype) for col in df.columns}

get_shape(df)

Get DataFrame shape.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required

Returns:

Type Description
tuple

(rows, columns)

Source code in odibi/engine/pandas_engine.py
def get_shape(self, df: Any) -> tuple:
    """Get DataFrame shape.

    Args:
        df: DataFrame or LazyDataset

    Returns:
        (rows, columns)
    """
    if isinstance(df, LazyDataset):
        cols = len(self.get_schema(df))
        rows = self.count_rows(df)
        return (rows, cols)
    return df.shape

get_source_files(df)

Get list of source files that generated this DataFrame.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset

required

Returns:

Type Description
List[str]

List of file paths

Source code in odibi/engine/pandas_engine.py
def get_source_files(self, df: Any) -> List[str]:
    """Get list of source files that generated this DataFrame.

    Args:
        df: DataFrame or LazyDataset

    Returns:
        List of file paths
    """
    if isinstance(df, LazyDataset):
        if isinstance(df.path, list):
            return df.path
        return [str(df.path)]

    if hasattr(df, "attrs"):
        return df.attrs.get("odibi_source_files", [])
    return []

get_table_schema(connection, table=None, path=None, format=None)

Get schema of an existing table/file.

Parameters:

Name Type Description Default
connection Any

Connection object providing base path and storage options.

required
table Optional[str]

Table name (for SQL sources).

None
path Optional[str]

File path (for file sources).

None
format Optional[str]

Data format (delta, parquet, csv, sql, sql_server, azure_sql).

None

Returns:

Type Description
Optional[Dict[str, str]]

Dict mapping column names to type strings, or None if schema

Optional[Dict[str, str]]

cannot be inferred.

Source code in odibi/engine/pandas_engine.py
def get_table_schema(
    self,
    connection: Any,
    table: Optional[str] = None,
    path: Optional[str] = None,
    format: Optional[str] = None,
) -> Optional[Dict[str, str]]:
    """Get schema of an existing table/file.

    Args:
        connection: Connection object providing base path and storage options.
        table: Table name (for SQL sources).
        path: File path (for file sources).
        format: Data format (delta, parquet, csv, sql, sql_server, azure_sql).

    Returns:
        Dict mapping column names to type strings, or None if schema
        cannot be inferred.
    """
    try:
        if table and is_sql_format(format):
            query = connection.build_select_query(table, limit=0)
            df = connection.read_sql(query)
            return self.get_schema(df)

        if path:
            full_path = connection.get_path(path)
            if not os.path.exists(full_path):
                return None

            if format == "delta":
                from deltalake import DeltaTable

                dt = DeltaTable(full_path)
                # Use pyarrow schema to pandas schema to avoid reading data
                arrow_schema = dt.schema().to_pyarrow()
                empty_df = arrow_schema.empty_table().to_pandas()
                return self.get_schema(empty_df)

            elif format == "parquet":
                import pyarrow.parquet as pq

                target_path = full_path
                if os.path.isdir(full_path):
                    # Find first parquet file
                    files = glob.glob(os.path.join(full_path, "*.parquet"))
                    if not files:
                        return None
                    target_path = files[0]

                schema = pq.read_schema(target_path)
                empty_df = schema.empty_table().to_pandas()
                return self.get_schema(empty_df)

            elif format == "csv":
                df = pd.read_csv(full_path, nrows=0)
                return self.get_schema(df)

    except (FileNotFoundError, PermissionError):
        return None
    except ImportError as e:
        # Log missing optional dependency
        import logging

        logging.getLogger(__name__).warning(
            f"Could not infer schema due to missing dependency: {e}"
        )
        return None
    except Exception as e:
        import logging

        logging.getLogger(__name__).warning(f"Failed to infer schema for {table or path}: {e}")
        return None
    return None

harmonize_schema(df, target_schema, policy)

Harmonize DataFrame schema with target schema according to policy.

Parameters:

Name Type Description Default
df DataFrame

Input DataFrame to harmonize.

required
target_schema Dict[str, str]

Target schema as dict of column_name -> type_string.

required
policy Any

Schema policy defining how to handle missing/new columns.

required

Returns:

Type Description
DataFrame

DataFrame with schema harmonized to match target_schema.

Raises:

Type Description
ValueError

If policy validation fails for missing or new columns.

Source code in odibi/engine/pandas_engine.py
def harmonize_schema(
    self, df: pd.DataFrame, target_schema: Dict[str, str], policy: Any
) -> pd.DataFrame:
    """Harmonize DataFrame schema with target schema according to policy.

    Args:
        df: Input DataFrame to harmonize.
        target_schema: Target schema as dict of column_name -> type_string.
        policy: Schema policy defining how to handle missing/new columns.

    Returns:
        DataFrame with schema harmonized to match target_schema.

    Raises:
        ValueError: If policy validation fails for missing or new columns.
    """
    # Ensure materialization
    df = self.materialize(df)

    from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

    target_cols = list(target_schema.keys())
    current_cols = df.columns.tolist()

    missing = set(target_cols) - set(current_cols)
    new_cols = set(current_cols) - set(target_cols)

    # 1. Check Validations
    if missing and policy.on_missing_columns == OnMissingColumns.FAIL:
        raise ValueError(f"Schema Policy Violation: Missing columns {missing}")

    if new_cols and policy.on_new_columns == OnNewColumns.FAIL:
        raise ValueError(f"Schema Policy Violation: New columns {new_cols}")

    # 2. Apply Transformations
    if policy.mode == SchemaMode.EVOLVE and policy.on_new_columns == OnNewColumns.ADD_NULLABLE:
        # Evolve: Add missing columns, Keep new columns
        for col in missing:
            df[col] = None
    else:
        # Enforce / Ignore New: Project to target schema (Drops new, Adds missing)
        # Note: reindex adds NaN for missing columns
        df = df.reindex(columns=target_cols)

    return df

maintain_table(connection, format, table=None, path=None, config=None)

Run table maintenance operations (optimize, vacuum).

Parameters:

Name Type Description Default
connection Any

Connection object providing base path and storage options.

required
format str

Data format (only delta is supported).

required
table Optional[str]

Table name (mutually exclusive with path).

None
path Optional[str]

File path (mutually exclusive with table).

None
config Optional[Any]

Maintenance configuration specifying operations to run.

None
Source code in odibi/engine/pandas_engine.py
def maintain_table(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    config: Optional[Any] = None,
) -> None:
    """Run table maintenance operations (optimize, vacuum).

    Args:
        connection: Connection object providing base path and storage options.
        format: Data format (only delta is supported).
        table: Table name (mutually exclusive with path).
        path: File path (mutually exclusive with table).
        config: Maintenance configuration specifying operations to run.
    """
    ctx = get_logging_context().with_context(engine="pandas")

    if format != "delta" or not config or not config.enabled:
        return

    if not path and not table:
        return

    full_path = connection.get_path(path if path else table)
    start = time.time()

    ctx.info("Starting table maintenance", path=str(full_path))

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.warning(
            "Auto-optimize skipped: 'deltalake' library not installed",
            path=str(full_path),
        )
        return

    try:
        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)

        ctx.info("Running Delta OPTIMIZE (compaction)", path=str(full_path))
        dt.optimize.compact()

        retention = config.vacuum_retention_hours
        if retention is not None and retention > 0:
            ctx.info(
                "Running Delta VACUUM",
                path=str(full_path),
                retention_hours=retention,
            )
            dt.vacuum(
                retention_hours=retention,
                enforce_retention_duration=True,
                dry_run=False,
            )

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Table maintenance completed",
            path=str(full_path),
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        ctx.warning(
            "Auto-optimize failed",
            path=str(full_path),
            error=str(e),
        )

materialize(df)

Materialize lazy dataset.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyDataset to materialize.

required

Returns:

Type Description
Any

Materialized pandas DataFrame. If input is already a DataFrame,

Any

returns it unchanged.

Source code in odibi/engine/pandas_engine.py
def materialize(self, df: Any) -> Any:
    """Materialize lazy dataset.

    Args:
        df: DataFrame or LazyDataset to materialize.

    Returns:
        Materialized pandas DataFrame. If input is already a DataFrame,
        returns it unchanged.
    """
    if isinstance(df, LazyDataset):
        # Re-invoke read but force materialization (by bypassing Lazy check)
        # We pass the resolved path directly
        # Note: We need to handle the case where path was resolved.
        # LazyDataset.path should be the FULL path.
        return self._read_file(
            full_path=df.path, format=df.format, options=df.options, connection=df.connection
        )
    return df

profile_nulls(df)

Calculate null percentage for each column.

Parameters:

Name Type Description Default
df DataFrame

DataFrame

required

Returns:

Type Description
Dict[str, float]

Dictionary of {column_name: null_percentage}

Source code in odibi/engine/pandas_engine.py
def profile_nulls(self, df: pd.DataFrame) -> Dict[str, float]:
    """Calculate null percentage for each column.

    Args:
        df: DataFrame

    Returns:
        Dictionary of {column_name: null_percentage}
    """
    # Ensure materialization
    df = self.materialize(df)

    # mean() of boolean DataFrame gives the percentage of True values
    return df.isna().mean().to_dict()

read(connection, format, table=None, path=None, streaming=False, schema=None, options=None, as_of_version=None, as_of_timestamp=None)

Read data using Pandas (or LazyDataset).

Parameters:

Name Type Description Default
connection Any

Connection object providing base path and storage options.

required
format str

Data format (csv, parquet, delta, json, excel, avro, sql).

required
table Optional[str]

Table name (mutually exclusive with path).

None
path Optional[str]

File path (mutually exclusive with table).

None
streaming bool

Whether to enable streaming mode (not supported in Pandas).

False
schema Optional[str]

Optional schema specification (not used in Pandas engine).

None
options Optional[Dict[str, Any]]

Additional read options to pass to pandas readers.

None
as_of_version Optional[int]

Delta Lake version for time travel queries.

None
as_of_timestamp Optional[str]

Delta Lake timestamp for time travel queries.

None

Returns:

Type Description
Union[DataFrame, Iterator[DataFrame]]

DataFrame or iterator of DataFrames containing the loaded data.

Raises:

Type Description
ValueError

If streaming is requested or neither path nor table is provided.

Source code in odibi/engine/pandas_engine.py
def read(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    streaming: bool = False,
    schema: Optional[str] = None,
    options: Optional[Dict[str, Any]] = None,
    as_of_version: Optional[int] = None,
    as_of_timestamp: Optional[str] = None,
) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
    """Read data using Pandas (or LazyDataset).

    Args:
        connection: Connection object providing base path and storage options.
        format: Data format (csv, parquet, delta, json, excel, avro, sql).
        table: Table name (mutually exclusive with path).
        path: File path (mutually exclusive with table).
        streaming: Whether to enable streaming mode (not supported in Pandas).
        schema: Optional schema specification (not used in Pandas engine).
        options: Additional read options to pass to pandas readers.
        as_of_version: Delta Lake version for time travel queries.
        as_of_timestamp: Delta Lake timestamp for time travel queries.

    Returns:
        DataFrame or iterator of DataFrames containing the loaded data.

    Raises:
        ValueError: If streaming is requested or neither path nor table is provided.
    """
    ctx = get_logging_context().with_context(engine="pandas")
    start = time.time()

    source = path or table
    ctx.debug(
        "Starting read operation",
        format=format,
        path=source,
        streaming=streaming,
        use_arrow=self.use_arrow,
    )

    if streaming:
        ctx.error(
            "Streaming not supported in Pandas engine",
            format=format,
            path=source,
        )
        raise ValueError(
            "Streaming is not supported in the Pandas engine. "
            "Please use 'engine: spark' for streaming pipelines."
        )

    options = options or {}

    # Handle simulation format (no path/table needed)
    if format == "simulation":
        ctx.debug("Generating simulated data")
        return self._read_simulation(options, ctx)

    # Resolve full path from connection
    try:
        full_path = self._resolve_path(path or table, connection)
    except ValueError:
        if table and not connection:
            ctx.error("Connection required when specifying 'table'", table=table)
            raise ValueError(
                f"Cannot read table '{table}': connection is required when using 'table' parameter. "
                "Provide a valid connection object or use 'path' for file-based reads."
            )
        ctx.error("Neither path nor table provided for read operation")
        raise ValueError(
            "Read operation failed: neither 'path' nor 'table' was provided. "
            "Specify a file path or table name in your configuration."
        )

    # Merge storage options for cloud connections
    merged_options = self._merge_storage_options(connection, options)

    # Sanitize options for pandas compatibility
    if "header" in merged_options:
        if merged_options["header"] is True:
            merged_options["header"] = 0
        elif merged_options["header"] is False:
            merged_options["header"] = None

    # Handle Time Travel options
    if as_of_version is not None:
        merged_options["versionAsOf"] = as_of_version
        ctx.debug("Time travel enabled", version=as_of_version)
    if as_of_timestamp is not None:
        merged_options["timestampAsOf"] = as_of_timestamp
        ctx.debug("Time travel enabled", timestamp=as_of_timestamp)

    # Check for Lazy/DuckDB optimization
    can_lazy_load = False

    if can_lazy_load:
        ctx.debug("Using lazy loading via DuckDB", path=str(full_path))
        if isinstance(full_path, (str, Path)):
            return LazyDataset(
                path=str(full_path),
                format=format,
                options=merged_options,
                connection=connection,
            )
        elif isinstance(full_path, list):
            return LazyDataset(
                path=full_path, format=format, options=merged_options, connection=connection
            )

    result = self._read_file(full_path, format, merged_options, connection)

    # Log metrics for materialized DataFrames
    elapsed = (time.time() - start) * 1000
    if isinstance(result, pd.DataFrame):
        row_count = len(result)
        memory_mb = result.memory_usage(deep=True).sum() / (1024 * 1024)

        ctx.log_file_io(
            path=str(full_path) if not isinstance(full_path, list) else str(full_path[0]),
            format=format,
            mode="read",
            rows=row_count,
        )
        ctx.log_pandas_metrics(
            memory_mb=memory_mb,
            dtypes={col: str(dtype) for col, dtype in result.dtypes.items()},
        )
        ctx.info(
            "Read completed",
            format=format,
            rows=row_count,
            elapsed_ms=round(elapsed, 2),
            memory_mb=round(memory_mb, 2),
        )

    return result

restore_delta(connection, path, version)

Restore Delta table to a specific version.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
version int

Version number to restore to

required
Source code in odibi/engine/pandas_engine.py
def restore_delta(self, connection: Any, path: str, version: int) -> None:
    """Restore Delta table to a specific version.

    Args:
        connection: Connection object
        path: Delta table path
        version: Version number to restore to
    """
    ctx = get_logging_context().with_context(engine="pandas")
    start = time.time()

    ctx.info("Starting Delta table restore", path=path, target_version=version)

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[pandas]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    dt.restore(version)

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta table restored",
        path=str(full_path),
        restored_to_version=version,
        elapsed_ms=round(elapsed, 2),
    )

table_exists(connection, table=None, path=None)

Check if table or location exists.

Parameters:

Name Type Description Default
connection Any

Connection object

required
table Optional[str]

Table name (not used in Pandas—no catalog)

None
path Optional[str]

File path

None

Returns:

Type Description
bool

True if file/directory exists, False otherwise

Source code in odibi/engine/pandas_engine.py
def table_exists(
    self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
) -> bool:
    """Check if table or location exists.

    Args:
        connection: Connection object
        table: Table name (not used in Pandas—no catalog)
        path: File path

    Returns:
        True if file/directory exists, False otherwise
    """
    if path:
        full_path = connection.get_path(path)
        return os.path.exists(full_path)
    return False

vacuum_delta(connection, path, retention_hours=168, dry_run=False, enforce_retention_duration=True)

VACUUM a Delta table to remove old files.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
retention_hours int

Retention period (default 168 = 7 days)

168
dry_run bool

If True, only show files to be deleted

False
enforce_retention_duration bool

If False, allows retention < 168 hours (testing only)

True

Returns:

Type Description
Dict[str, Any]

Dictionary with files_deleted count

Source code in odibi/engine/pandas_engine.py
def vacuum_delta(
    self,
    connection: Any,
    path: str,
    retention_hours: int = 168,
    dry_run: bool = False,
    enforce_retention_duration: bool = True,
) -> Dict[str, Any]:
    """VACUUM a Delta table to remove old files.

    Args:
        connection: Connection object
        path: Delta table path
        retention_hours: Retention period (default 168 = 7 days)
        dry_run: If True, only show files to be deleted
        enforce_retention_duration: If False, allows retention < 168 hours (testing only)

    Returns:
        Dictionary with files_deleted count
    """
    ctx = get_logging_context().with_context(engine="pandas")
    start = time.time()

    ctx.debug(
        "Starting Delta VACUUM",
        path=path,
        retention_hours=retention_hours,
        dry_run=dry_run,
    )

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[pandas]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    deleted_files = dt.vacuum(
        retention_hours=retention_hours,
        dry_run=dry_run,
        enforce_retention_duration=enforce_retention_duration,
    )

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta VACUUM completed",
        path=str(full_path),
        files_deleted=len(deleted_files),
        dry_run=dry_run,
        elapsed_ms=round(elapsed, 2),
    )

    return {"files_deleted": len(deleted_files)}

validate_data(df, validation_config)

Validate DataFrame against rules.

Parameters:

Name Type Description Default
df DataFrame

DataFrame

required
validation_config Any

ValidationConfig object

required

Returns:

Type Description
List[str]

List of validation failure messages

Source code in odibi/engine/pandas_engine.py
def validate_data(self, df: pd.DataFrame, validation_config: Any) -> List[str]:
    """Validate DataFrame against rules.

    Args:
        df: DataFrame
        validation_config: ValidationConfig object

    Returns:
        List of validation failure messages
    """
    # Ensure materialization
    df = self.materialize(df)

    failures = []

    # Check not empty
    if validation_config.not_empty:
        if len(df) == 0:
            failures.append("DataFrame is empty")

    # Check for nulls in specified columns
    if validation_config.no_nulls:
        null_counts = self.count_nulls(df, validation_config.no_nulls)
        for col, count in null_counts.items():
            if count > 0:
                failures.append(f"Column '{col}' has {count} null values")

    # Schema validation
    if validation_config.schema_validation:
        schema_failures = self.validate_schema(df, validation_config.schema_validation)
        failures.extend(schema_failures)

    # Range validation
    if validation_config.ranges:
        for col, bounds in validation_config.ranges.items():
            if col in df.columns:
                min_val = bounds.get("min")
                max_val = bounds.get("max")

                if min_val is not None:
                    min_violations = df[df[col] < min_val]
                    if len(min_violations) > 0:
                        failures.append(f"Column '{col}' has values < {min_val}")

                if max_val is not None:
                    max_violations = df[df[col] > max_val]
                    if len(max_violations) > 0:
                        failures.append(f"Column '{col}' has values > {max_val}")
            else:
                failures.append(f"Column '{col}' not found for range validation")

    # Allowed values validation
    if validation_config.allowed_values:
        for col, allowed in validation_config.allowed_values.items():
            if col in df.columns:
                # Check for values not in allowed list
                invalid = df[~df[col].isin(allowed)]
                if len(invalid) > 0:
                    failures.append(f"Column '{col}' has invalid values")
            else:
                failures.append(f"Column '{col}' not found for allowed values validation")

    return failures

validate_schema(df, schema_rules)

Validate DataFrame schema.

Parameters:

Name Type Description Default
df DataFrame

DataFrame

required
schema_rules Dict[str, Any]

Validation rules

required

Returns:

Type Description
List[str]

List of validation failures

Source code in odibi/engine/pandas_engine.py
def validate_schema(self, df: pd.DataFrame, schema_rules: Dict[str, Any]) -> List[str]:
    """Validate DataFrame schema.

    Args:
        df: DataFrame
        schema_rules: Validation rules

    Returns:
        List of validation failures
    """
    # Ensure materialization
    df = self.materialize(df)

    failures = []

    # Check required columns
    if "required_columns" in schema_rules:
        required = schema_rules["required_columns"]
        missing = set(required) - set(df.columns)
        if missing:
            failures.append(f"Missing required columns: {', '.join(missing)}")

    # Check column types
    if "types" in schema_rules:
        type_map = {
            "int": ["int64", "int32", "int16", "int8"],
            "float": ["float64", "float32"],
            "str": ["object", "string"],
            "bool": ["bool"],
        }

        for col, expected_type in schema_rules["types"].items():
            if col not in df.columns:
                failures.append(f"Column '{col}' not found for type validation")
                continue

            actual_type = str(df[col].dtype)
            # Handle pyarrow types (e.g. int64[pyarrow])
            if "[" in actual_type and "pyarrow" in actual_type:
                actual_type = actual_type.split("[")[0]

            expected_dtypes = type_map.get(expected_type, [expected_type])

            if actual_type not in expected_dtypes:
                failures.append(
                    f"Column '{col}' has type '{actual_type}', expected '{expected_type}'"
                )

    return failures

write(df, connection, format, table=None, path=None, register_table=None, mode='overwrite', options=None, streaming_config=None)

Write data using Pandas.

Parameters:

Name Type Description Default
df Union[DataFrame, Iterator[DataFrame]]

DataFrame or iterator of DataFrames to write.

required
connection Any

Connection object providing base path and storage options.

required
format str

Output format (csv, parquet, delta, json, excel, avro, sql).

required
table Optional[str]

Table name (mutually exclusive with path).

None
path Optional[str]

File path (mutually exclusive with table).

None
register_table Optional[str]

Optional table name to register in catalog (not used).

None
mode str

Write mode (overwrite, append, upsert, append_once).

'overwrite'
options Optional[Dict[str, Any]]

Additional write options to pass to pandas writers.

None
streaming_config Optional[Any]

Streaming configuration (not used in Pandas engine).

None

Returns:

Type Description
Optional[Dict[str, Any]]

Optional dict with write statistics for Delta writes, None otherwise.

Raises:

Type Description
ValueError

If neither path nor table is provided.

Source code in odibi/engine/pandas_engine.py
def write(
    self,
    df: Union[pd.DataFrame, Iterator[pd.DataFrame]],
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    register_table: Optional[str] = None,
    mode: str = "overwrite",
    options: Optional[Dict[str, Any]] = None,
    streaming_config: Optional[Any] = None,
) -> Optional[Dict[str, Any]]:
    """Write data using Pandas.

    Args:
        df: DataFrame or iterator of DataFrames to write.
        connection: Connection object providing base path and storage options.
        format: Output format (csv, parquet, delta, json, excel, avro, sql).
        table: Table name (mutually exclusive with path).
        path: File path (mutually exclusive with table).
        register_table: Optional table name to register in catalog (not used).
        mode: Write mode (overwrite, append, upsert, append_once).
        options: Additional write options to pass to pandas writers.
        streaming_config: Streaming configuration (not used in Pandas engine).

    Returns:
        Optional dict with write statistics for Delta writes, None otherwise.

    Raises:
        ValueError: If neither path nor table is provided.
    """
    ctx = get_logging_context().with_context(engine="pandas")
    start = time.time()

    destination = path or table
    ctx.debug(
        "Starting write operation",
        format=format,
        destination=destination,
        mode=mode,
    )

    # Ensure materialization if LazyDataset
    df = self.materialize(df)

    options = options or {}

    # Handle iterator/generator input
    from collections.abc import Iterator

    if isinstance(df, Iterator):
        ctx.debug("Writing iterator/generator input")
        return self._write_iterator(df, connection, format, table, path, mode, options)

    row_count = len(df)
    memory_mb = df.memory_usage(deep=True).sum() / (1024 * 1024)

    ctx.log_pandas_metrics(
        memory_mb=memory_mb,
        dtypes={col: str(dtype) for col, dtype in df.dtypes.items()},
    )

    # SQL Server / Azure SQL Support
    if is_sql_format(format):
        ctx.debug("Writing to SQL", table=table, mode=mode)
        return self._write_sql(df, connection, table, mode, options)

    # Resolve full path from connection
    try:
        full_path = self._resolve_path(path or table, connection)
    except ValueError:
        if table and not connection:
            ctx.error("Connection required when specifying 'table'", table=table)
            raise ValueError("Connection is required when specifying 'table'.")
        ctx.error("Neither path nor table provided for write operation")
        raise ValueError("Either path or table must be provided")

    # Merge storage options for cloud connections
    merged_options = self._merge_storage_options(connection, options)

    # Custom Writers
    if format in self._custom_writers:
        ctx.debug(f"Using custom writer for format: {format}")
        writer_options = merged_options.copy()
        writer_options.pop("keys", None)
        self._custom_writers[format](df, full_path, mode=mode, **writer_options)
        return None

    # Ensure directory exists (local only)
    self._ensure_directory(full_path)

    # Warn about partitioning
    self._check_partitioning(merged_options)

    # Delta Lake Write
    if format == "delta":
        ctx.debug("Writing Delta table", path=str(full_path), mode=mode)
        result = self._write_delta(df, full_path, mode, merged_options)
        elapsed = (time.time() - start) * 1000
        ctx.log_file_io(
            path=str(full_path),
            format=format,
            mode=mode,
            rows=row_count,
        )
        ctx.info(
            "Write completed",
            format=format,
            rows=row_count,
            elapsed_ms=round(elapsed, 2),
        )
        return result

    # Handle Generic Upsert/Append-Once for non-Delta
    if mode in ["upsert", "append_once"]:
        ctx.debug(f"Handling {mode} mode for non-Delta format")
        df, mode = self._handle_generic_upsert(df, full_path, format, mode, merged_options)
        row_count = len(df)

    # Standard File Write
    result = self._write_file(df, full_path, format, mode, merged_options)

    elapsed = (time.time() - start) * 1000
    ctx.log_file_io(
        path=str(full_path),
        format=format,
        mode=mode,
        rows=row_count,
    )
    ctx.info(
        "Write completed",
        format=format,
        rows=row_count,
        elapsed_ms=round(elapsed, 2),
    )

    return result

odibi.engine.spark_engine

Spark execution engine (Phase 2B: Delta Lake support).

Status: Phase 2B implemented - Delta Lake read/write, VACUUM, history, restore

SparkEngine

Bases: Engine

Spark execution engine with PySpark backend.

Phase 2A: Basic read/write + ADLS multi-account support Phase 2B: Delta Lake support

Source code in odibi/engine/spark_engine.py
  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
 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
 433
 434
 435
 436
 437
 438
 439
 440
 441
 442
 443
 444
 445
 446
 447
 448
 449
 450
 451
 452
 453
 454
 455
 456
 457
 458
 459
 460
 461
 462
 463
 464
 465
 466
 467
 468
 469
 470
 471
 472
 473
 474
 475
 476
 477
 478
 479
 480
 481
 482
 483
 484
 485
 486
 487
 488
 489
 490
 491
 492
 493
 494
 495
 496
 497
 498
 499
 500
 501
 502
 503
 504
 505
 506
 507
 508
 509
 510
 511
 512
 513
 514
 515
 516
 517
 518
 519
 520
 521
 522
 523
 524
 525
 526
 527
 528
 529
 530
 531
 532
 533
 534
 535
 536
 537
 538
 539
 540
 541
 542
 543
 544
 545
 546
 547
 548
 549
 550
 551
 552
 553
 554
 555
 556
 557
 558
 559
 560
 561
 562
 563
 564
 565
 566
 567
 568
 569
 570
 571
 572
 573
 574
 575
 576
 577
 578
 579
 580
 581
 582
 583
 584
 585
 586
 587
 588
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
class SparkEngine(Engine):
    """Spark execution engine with PySpark backend.

    Phase 2A: Basic read/write + ADLS multi-account support
    Phase 2B: Delta Lake support
    """

    name = "spark"
    engine_type = EngineType.SPARK

    def __init__(
        self,
        connections: Optional[Dict[str, Any]] = None,
        spark_session: Any = None,
        config: Optional[Dict[str, Any]] = None,
    ):
        """Initialize Spark engine with import guard.

        Args:
            connections: Dictionary of connection objects (for multi-account config)
            spark_session: Existing SparkSession (optional, creates new if None)
            config: Engine configuration (optional)

        Raises:
            ImportError: If pyspark not installed
        """
        ctx = get_logging_context().with_context(engine="spark")
        ctx.debug("Initializing SparkEngine", connections_count=len(connections or {}))

        try:
            from pyspark.sql import SparkSession
        except ImportError as e:
            ctx.error(
                "PySpark not installed",
                error_type="ImportError",
                suggestion="pip install odibi[spark]",
            )
            raise ImportError(
                "Spark support requires 'pip install odibi[spark]'. "
                "See docs/setup_databricks.md for setup instructions."
            ) from e

        start_time = time.time()

        _reused_active = False
        if spark_session:
            self.spark = spark_session
        elif SparkSession.getActiveSession() is not None:
            # Reuse the SparkSession already running in this process. Every
            # notebook context has one — Databricks serverless and classic
            # clusters, Jupyter+Spark, an existing Spark Connect client. Building
            # a fresh session here instead would spin up a local Spark or a new
            # Connect client with no sc:// URL, which is the INVALID_CONNECT_URL
            # failure on serverless (and configure_spark_with_delta_pip assumes
            # local Spark). Only build below when there is genuinely no session.
            self.spark = SparkSession.getActiveSession()
            _reused_active = True
            ctx.debug("Reusing active SparkSession (no new session created)")
        else:
            # Configure Delta Lake support
            try:
                from delta import configure_spark_with_delta_pip

                builder = SparkSession.builder.appName("odibi").config(
                    "spark.sql.sources.partitionOverwriteMode", "dynamic"
                )

                # Performance Optimizations
                builder = builder.config("spark.sql.execution.arrow.pyspark.enabled", "true")
                builder = builder.config("spark.sql.adaptive.enabled", "true")

                # Reduce Verbosity
                builder = builder.config(
                    "spark.driver.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
                )
                builder = builder.config(
                    "spark.executor.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
                )

                self.spark = configure_spark_with_delta_pip(builder).getOrCreate()
                try:
                    self.spark.sparkContext.setLogLevel("ERROR")
                except Exception:
                    pass  # Spark Connect: sparkContext not available

                ctx.debug("Delta Lake support enabled")

            except ImportError:
                ctx.debug("Delta Lake not available, using standard Spark")
                builder = SparkSession.builder.appName("odibi").config(
                    "spark.sql.sources.partitionOverwriteMode", "dynamic"
                )

                # Performance Optimizations
                builder = builder.config("spark.sql.execution.arrow.pyspark.enabled", "true")
                builder = builder.config("spark.sql.adaptive.enabled", "true")

                # Reduce Verbosity
                builder = builder.config(
                    "spark.driver.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
                )

                self.spark = builder.getOrCreate()
                try:
                    self.spark.sparkContext.setLogLevel("ERROR")
                except Exception:
                    pass  # Spark Connect: sparkContext not available

        self.config = config or {}
        self.connections = connections or {}

        # Configure all ADLS connections upfront
        self._configure_all_connections()

        # Apply user-defined Spark configs from performance settings
        self._apply_spark_config()

        elapsed = (time.time() - start_time) * 1000
        try:
            app_name = self.spark.sparkContext.appName
        except Exception:
            app_name = "unknown"  # Spark Connect: sparkContext not available
        ctx.info(
            "SparkEngine initialized",
            elapsed_ms=round(elapsed, 2),
            app_name=app_name,
            spark_version=self.spark.version,
            connections_configured=len(self.connections),
            using_existing_session=spark_session is not None or _reused_active,
        )

    def _configure_all_connections(self) -> None:
        """Configure Spark with all ADLS connection credentials.

        This sets all storage account keys upfront so Spark can access
        multiple accounts. Keys are scoped by account name, so no conflicts.
        """
        ctx = get_logging_context().with_context(engine="spark")

        for conn_name, connection in self.connections.items():
            if hasattr(connection, "configure_spark"):
                ctx.log_connection(
                    connection_type=type(connection).__name__,
                    connection_name=conn_name,
                    action="configure_spark",
                )
                try:
                    connection.configure_spark(self.spark)
                    ctx.debug(f"Configured ADLS connection: {conn_name}")
                except Exception as e:
                    ctx.error(
                        f"Failed to configure ADLS connection: {conn_name}",
                        error_type=type(e).__name__,
                        error_message=str(e),
                    )
                    raise

    def _apply_spark_config(self) -> None:
        """Apply user-defined Spark configurations from performance settings.

        Applies configs via spark.conf.set() for runtime-settable options.
        For existing sessions (e.g., Databricks), only modifiable configs take effect.
        """
        ctx = get_logging_context().with_context(engine="spark")

        performance = self.config.get("performance", {})
        spark_config = performance.get("spark_config", {})

        if not spark_config:
            return

        ctx.debug("Applying Spark configuration", config_count=len(spark_config))

        for key, value in spark_config.items():
            try:
                self.spark.conf.set(key, value)
                ctx.debug(
                    f"Applied Spark config: {key}={value}", config_key=key, config_value=value
                )
            except Exception as e:
                ctx.warning(
                    f"Failed to set Spark config '{key}'",
                    config_key=key,
                    error_message=str(e),
                    suggestion="This config may require session restart",
                )

    def _safe_partition_count(self, df) -> int:
        """Get DataFrame partition count, with fallback for Spark Connect.

        On shared/user-isolation clusters (Spark Connect), df.rdd is blocked.
        Since partition count is only used for logging/metrics, we return -1
        as a safe fallback rather than crashing the pipeline.
        """
        try:
            return df.rdd.getNumPartitions()
        except Exception:
            return -1

    def _apply_table_properties(
        self, target: str, properties: Dict[str, str], is_table: bool = False
    ) -> None:
        """Apply table properties to a Delta table.

        Performance: Batches all properties into a single ALTER TABLE statement
        to avoid multiple round-trips to the catalog.
        """
        if not properties:
            return

        ctx = get_logging_context().with_context(engine="spark")

        try:
            table_ref = target if is_table else f"delta.`{target}`"
            ctx.debug(
                f"Applying table properties to {target}",
                properties_count=len(properties),
                is_table=is_table,
            )

            props_list = [f"'{k}' = '{v}'" for k, v in properties.items()]
            props_str = ", ".join(props_list)
            sql = f"ALTER TABLE {table_ref} SET TBLPROPERTIES ({props_str})"
            self.spark.sql(sql)
            ctx.debug(f"Set {len(properties)} table properties in single statement")

        except Exception as e:
            ctx.warning(
                f"Failed to set table properties on {target}",
                error_type=type(e).__name__,
                error_message=str(e),
            )

    def _optimize_delta_write(
        self, target: str, options: Dict[str, Any], is_table: bool = False
    ) -> None:
        """Run Delta Lake optimization (OPTIMIZE / ZORDER)."""
        should_optimize = options.get("optimize_write", False)
        zorder_by = options.get("zorder_by")

        if not should_optimize and not zorder_by:
            return

        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        try:
            if is_table:
                sql = f"OPTIMIZE {target}"
            else:
                sql = f"OPTIMIZE delta.`{target}`"

            if zorder_by:
                if isinstance(zorder_by, str):
                    zorder_by = [zorder_by]
                cols = ", ".join(zorder_by)
                sql += f" ZORDER BY ({cols})"

            ctx.debug("Running Delta optimization", sql=sql, target=target)
            self.spark.sql(sql)

            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                "Delta optimization completed",
                target=target,
                zorder_by=zorder_by,
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.warning(
                f"Optimization failed for {target}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )

    def _get_last_delta_commit_info(
        self, target: str, is_table: bool = False
    ) -> Optional[Dict[str, Any]]:
        """Get metadata for the most recent Delta commit."""
        ctx = get_logging_context().with_context(engine="spark")

        try:
            from delta.tables import DeltaTable

            if is_table:
                dt = DeltaTable.forName(self.spark, target)
            else:
                dt = DeltaTable.forPath(self.spark, target)

            last_commit = dt.history(1).collect()[0]

            def safe_get(row, field):
                if hasattr(row, field):
                    return getattr(row, field)
                if hasattr(row, "__getitem__"):
                    try:
                        return row[field]
                    except (KeyError, ValueError):
                        return None
                return None

            commit_info = {
                "version": safe_get(last_commit, "version"),
                "timestamp": safe_get(last_commit, "timestamp"),
                "operation": safe_get(last_commit, "operation"),
                "operation_metrics": safe_get(last_commit, "operationMetrics"),
                "read_version": safe_get(last_commit, "readVersion"),
            }

            ctx.debug(
                "Delta commit metadata retrieved",
                target=target,
                version=commit_info.get("version"),
                operation=commit_info.get("operation"),
            )

            return commit_info

        except Exception as e:
            ctx.warning(
                f"Failed to fetch Delta commit info for {target}",
                error_type=type(e).__name__,
                error_message=str(e),
            )
            return None

    def harmonize_schema(self, df, target_schema: Dict[str, str], policy: Any):
        """Harmonize DataFrame schema with target schema according to policy."""
        from pyspark.sql.functions import col, lit

        from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

        ctx = get_logging_context().with_context(engine="spark")

        target_cols = list(target_schema.keys())
        current_cols = df.columns

        missing = set(target_cols) - set(current_cols)
        new_cols = set(current_cols) - set(target_cols)

        ctx.debug(
            "Schema harmonization",
            target_columns=len(target_cols),
            current_columns=len(current_cols),
            missing_columns=list(missing) if missing else None,
            new_columns=list(new_cols) if new_cols else None,
            policy_mode=policy.mode.value if hasattr(policy.mode, "value") else str(policy.mode),
        )

        # Check Validations
        if missing and policy.on_missing_columns == OnMissingColumns.FAIL:
            ctx.error(
                f"Schema Policy Violation: Missing columns {missing}",
                missing_columns=list(missing),
            )
            raise ValueError(f"Schema Policy Violation: Missing columns {missing}")

        if new_cols and policy.on_new_columns == OnNewColumns.FAIL:
            ctx.error(
                f"Schema Policy Violation: New columns {new_cols}",
                new_columns=list(new_cols),
            )
            raise ValueError(f"Schema Policy Violation: New columns {new_cols}")

        # Apply Transformations
        if policy.mode == SchemaMode.EVOLVE and policy.on_new_columns == OnNewColumns.ADD_NULLABLE:
            res = df
            for c in missing:
                res = res.withColumn(c, lit(None))
            ctx.debug("Schema evolved: added missing columns as null")
            return res
        else:
            select_exprs = []
            for c in target_cols:
                if c in current_cols:
                    # Escape column names with backticks to handle special characters
                    select_exprs.append(col(f"`{c}`"))
                else:
                    select_exprs.append(lit(None).alias(c))

            ctx.debug("Schema enforced: projected to target schema")
            return df.select(*select_exprs)

    def anonymize(self, df, columns: List[str], method: str, salt: Optional[str] = None):
        """Anonymize columns using Spark functions."""
        from pyspark.sql.functions import col, concat, lit, regexp_replace, sha2

        ctx = get_logging_context().with_context(engine="spark")
        ctx.debug(
            "Anonymizing columns",
            columns=columns,
            method=method,
            has_salt=salt is not None,
        )

        res = df
        for c in columns:
            if c not in df.columns:
                ctx.warning(f"Column '{c}' not found for anonymization, skipping", column=c)
                continue

            # Escape column names with backticks to handle special characters
            escaped_col = f"`{c}`"
            if method == "hash":
                if salt:
                    res = res.withColumn(c, sha2(concat(col(escaped_col), lit(salt)), 256))
                else:
                    res = res.withColumn(c, sha2(col(escaped_col), 256))

            elif method == "mask":
                res = res.withColumn(c, regexp_replace(col(escaped_col), ".(?=.{4})", "*"))

            elif method == "redact":
                res = res.withColumn(c, lit("[REDACTED]"))

        ctx.debug(f"Anonymization completed using {method}")
        return res

    def get_schema(self, df) -> Dict[str, str]:
        """Get DataFrame schema with types."""
        return {f.name: f.dataType.simpleString() for f in df.schema}

    def get_shape(self, df) -> Tuple[int, int]:
        """Get DataFrame shape as (rows, columns)."""
        return (df.count(), len(df.columns))

    def count_rows(self, df) -> int:
        """Count rows in DataFrame."""
        return df.count()

    def read(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        streaming: bool = False,
        schema: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        as_of_version: Optional[int] = None,
        as_of_timestamp: Optional[str] = None,
    ) -> Any:
        """Read data using Spark.

        Args:
            connection: Connection object (with get_path method)
            format: Data format (csv, parquet, json, delta, sql_server)
            table: Table name
            path: File path
            streaming: Whether to read as a stream (readStream)
            schema: Schema string in DDL format (required for streaming file sources)
            options: Format-specific options (including versionAsOf for Delta time travel)
            as_of_version: Time travel version
            as_of_timestamp: Time travel timestamp

        Returns:
            Spark DataFrame (or Streaming DataFrame)
        """
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()
        options = options or {}

        source_identifier = table or path or "unknown"
        ctx.debug(
            "Starting Spark read",
            format=format,
            source=source_identifier,
            streaming=streaming,
            as_of_version=as_of_version,
            as_of_timestamp=as_of_timestamp,
        )

        # Handle Time Travel options
        if as_of_version is not None:
            options["versionAsOf"] = as_of_version
            ctx.debug(f"Time travel enabled: version {as_of_version}")
        if as_of_timestamp is not None:
            options["timestampAsOf"] = as_of_timestamp
            ctx.debug(f"Time travel enabled: timestamp {as_of_timestamp}")

        # SQL Server / Azure SQL Support
        if is_sql_format(format):
            if streaming:
                ctx.error("Streaming not supported for SQL Server / Azure SQL")
                raise ValueError("Streaming not supported for SQL Server / Azure SQL yet.")

            if not hasattr(connection, "get_spark_options"):
                conn_type = type(connection).__name__
                msg = f"Connection type '{conn_type}' does not support Spark SQL read"
                ctx.error(msg, connection_type=conn_type)
                raise ValueError(msg)

            jdbc_options = connection.get_spark_options()
            merged_options = {**jdbc_options, **options}

            # Extract filter for SQL pushdown
            sql_filter = merged_options.pop("filter", None)

            if "query" in merged_options:
                merged_options.pop("dbtable", None)
                # If filter provided with query, append to WHERE clause
                if sql_filter:
                    existing_query = merged_options["query"]
                    # Wrap existing query and add filter
                    if "WHERE" in existing_query.upper():
                        merged_options["query"] = f"({existing_query}) AND ({sql_filter})"
                    else:
                        subquery = f"SELECT * FROM ({existing_query}) AS _subq WHERE {sql_filter}"
                        merged_options["query"] = subquery
                    ctx.debug(f"Applied SQL pushdown filter to query: {sql_filter}")
            elif table:
                # Build query with filter pushdown instead of using dbtable
                if sql_filter:
                    merged_options.pop("dbtable", None)
                    merged_options["query"] = f"SELECT * FROM {table} WHERE {sql_filter}"
                    ctx.debug(f"Applied SQL pushdown filter: {sql_filter}")
                else:
                    merged_options["dbtable"] = table
            elif "dbtable" not in merged_options:
                ctx.error("SQL format requires 'table' config or 'query' option")
                raise ValueError("SQL format requires 'table' config or 'query' option")

            ctx.debug("Executing JDBC read", has_query="query" in merged_options)

            try:
                df = self.spark.read.format("jdbc").options(**merged_options).load()
                elapsed = (time.time() - start_time) * 1000
                partition_count = self._safe_partition_count(df)

                ctx.log_file_io(path=source_identifier, format=format, mode="read")
                ctx.log_spark_metrics(partition_count=partition_count)
                ctx.info(
                    "JDBC read completed",
                    source=source_identifier,
                    elapsed_ms=round(elapsed, 2),
                    partitions=partition_count,
                )
                return df

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "JDBC read failed",
                    source=source_identifier,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        # Simulation format: delegate to Pandas, convert to Spark
        if format == "simulation":
            if streaming:
                ctx.error("Streaming not supported for simulation format")
                raise ValueError("Streaming not supported for simulation format")

            ctx.debug("Reading simulation via Pandas engine, converting to Spark")
            from odibi.engine.pandas_engine import PandasEngine

            # Use Pandas engine for simulation
            pandas_engine = PandasEngine()
            pdf = pandas_engine._read_simulation(options, ctx)

            # Convert Pandas DataFrame to Spark DataFrame
            df = self.spark.createDataFrame(pdf)
            elapsed = (time.time() - start_time) * 1000

            # Preserve HWM metadata in Spark DataFrame metadata
            if hasattr(pdf, "attrs") and "_simulation_max_timestamp" in pdf.attrs:
                # Note: Spark doesn't have attrs, so we'll attach as a property
                df._simulation_max_timestamp = pdf.attrs["_simulation_max_timestamp"]

            # Preserve random walk state
            if hasattr(pdf, "attrs") and "_simulation_random_walk_state" in pdf.attrs:
                df._simulation_random_walk_state = pdf.attrs["_simulation_random_walk_state"]

            # Preserve scheduled event state
            if hasattr(pdf, "attrs") and "_simulation_scheduled_event_state" in pdf.attrs:
                df._simulation_scheduled_event_state = pdf.attrs[
                    "_simulation_scheduled_event_state"
                ]

            ctx.info(
                "Simulation read completed (via Pandas)",
                elapsed_ms=round(elapsed, 2),
                row_count=len(pdf),
            )
            return df

        # Read based on format
        if table:
            # Managed/External Table (Catalog)
            ctx.debug(f"Reading from catalog table: {table}")

            if streaming:
                reader = self.spark.readStream.format(format)
            else:
                reader = self.spark.read.format(format)

            for key, value in options.items():
                reader = reader.option(key, value)

            try:
                df = reader.table(table)

                if "filter" in options:
                    df = df.filter(options["filter"])
                    ctx.debug(f"Applied filter: {options['filter']}")

                elapsed = (time.time() - start_time) * 1000

                if not streaming:
                    partition_count = self._safe_partition_count(df)
                    ctx.log_spark_metrics(partition_count=partition_count)
                    ctx.log_file_io(path=table, format=format, mode="read")
                    ctx.info(
                        f"Table read completed: {table}",
                        elapsed_ms=round(elapsed, 2),
                        partitions=partition_count,
                    )
                else:
                    ctx.info(f"Streaming read started: {table}", elapsed_ms=round(elapsed, 2))

                return df

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    f"Table read failed: {table}",
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        elif path:
            # File Path
            full_path = connection.get_path(path)
            ctx.debug(f"Reading from path: {full_path}")

            # Excel format: delegate to Pandas, convert to Spark
            if format == "excel":
                if streaming:
                    ctx.error("Streaming not supported for Excel format")
                    raise ValueError("Streaming not supported for Excel format")

                ctx.debug("Reading Excel via Pandas engine (best Excel support)")
                from odibi.engine.pandas_engine import PandasEngine

                # Get storage_options from connection for Azure/cloud authentication
                storage_options = None
                if hasattr(connection, "pandas_storage_options"):
                    storage_options = connection.pandas_storage_options()

                # Use Pandas engine for Excel reading
                pandas_engine = PandasEngine()
                pdf = pandas_engine._read_excel_with_patterns(
                    full_path,
                    sheet_pattern=options.pop("sheet_pattern", None),
                    sheet_pattern_case_sensitive=options.pop("sheet_pattern_case_sensitive", False),
                    add_source_file=options.pop("add_source_file", False),
                    is_glob="*" in str(full_path) or "?" in str(full_path),
                    ctx=ctx,
                    storage_options=storage_options,
                    **options,
                )

                # Convert Pandas DataFrame to Spark DataFrame
                df = self.spark.createDataFrame(pdf)
                elapsed = (time.time() - start_time) * 1000
                ctx.info(
                    f"Excel read completed (via Pandas): {path}",
                    elapsed_ms=round(elapsed, 2),
                    row_count=len(pdf),
                )
                return df

            # API format: delegate to ApiFetcher, convert to Spark
            if format == "api":
                if streaming:
                    ctx.error("Streaming not supported for API format")
                    raise ValueError("Streaming not supported for API format")

                from odibi.connections.api_fetcher import create_api_fetcher
                from odibi.connections.http import HttpConnection

                ctx.debug("Reading API via ApiFetcher", endpoint=str(path))

                if not isinstance(connection, HttpConnection):
                    raise ValueError(
                        f"Cannot read API data: connection type '{type(connection).__name__}' "
                        "is not an HttpConnection. Use an HTTP connection for API format."
                    )

                # Extract API-specific options
                api_options = options.copy()
                params = api_options.pop("params", {})
                max_records = api_options.pop("max_records", None)
                method = api_options.pop("method", "GET")
                request_body = api_options.pop("request_body", None)

                # Create fetcher with connection's base_url and headers
                fetcher = create_api_fetcher(
                    base_url=connection.base_url,
                    headers=connection.headers,
                    options=api_options,
                )

                # Fetch data as Pandas DataFrame
                pdf = fetcher.fetch_dataframe(
                    endpoint=str(full_path),
                    params=params,
                    max_records=max_records,
                    method=method,
                    request_body=request_body,
                )

                # Convert Pandas DataFrame to Spark DataFrame
                df = self.spark.createDataFrame(pdf)
                elapsed = (time.time() - start_time) * 1000
                ctx.info(
                    f"API read completed: {path}",
                    elapsed_ms=round(elapsed, 2),
                    row_count=len(pdf),
                )
                return df

            # Auto-detect encoding for CSV (Batch only)
            if not streaming and format == "csv" and options.get("auto_encoding"):
                options = options.copy()
                options.pop("auto_encoding")

                if "encoding" not in options:
                    try:
                        from odibi.utils.encoding import detect_encoding

                        detected = detect_encoding(connection, path)
                        if detected:
                            options["encoding"] = detected
                            ctx.debug(f"Detected encoding: {detected}", path=path)
                    except ImportError:
                        pass
                    except Exception as e:
                        ctx.warning(
                            f"Encoding detection failed for {path}",
                            error_message=str(e),
                        )

            # Resolve cloudFiles.schemaLocation through read connection if still relative
            # (Node level resolves via write connection first if available)
            if "cloudFiles.schemaLocation" in options:
                schema_location = options["cloudFiles.schemaLocation"]
                if not schema_location.startswith(
                    ("abfss://", "s3://", "gs://", "dbfs://", "hdfs://", "wasbs://")
                ):
                    options = options.copy()
                    options["cloudFiles.schemaLocation"] = connection.get_path(schema_location)
                    ctx.debug(
                        "Resolved cloudFiles.schemaLocation through read connection (fallback)",
                        original=schema_location,
                        resolved=options["cloudFiles.schemaLocation"],
                    )

            if streaming:
                reader = self.spark.readStream.format(format)
                if schema:
                    reader = reader.schema(schema)
                    ctx.debug(f"Applied schema for streaming read: {schema[:100]}...")
                else:
                    # Determine if we should warn about missing schema
                    # Formats that can infer schema: delta, parquet, avro (embedded schema)
                    # cloudFiles with schemaLocation or self-describing formats (avro, parquet) are fine
                    should_warn = True

                    if format in ["delta", "parquet"]:
                        should_warn = False
                    elif format == "cloudFiles":
                        cloud_format = options.get("cloudFiles.format", "")
                        has_schema_location = "cloudFiles.schemaLocation" in options
                        # avro and parquet have embedded schemas
                        if cloud_format in ["avro", "parquet"] or has_schema_location:
                            should_warn = False

                    if should_warn:
                        ctx.warning(
                            f"Streaming read from '{format}' format without schema. "
                            "Schema inference is not supported for streaming sources. "
                            "Consider adding 'schema' to your read config."
                        )
            else:
                reader = self.spark.read.format(format)
                if schema:
                    reader = reader.schema(schema)

            for key, value in options.items():
                if key == "header" and isinstance(value, bool):
                    value = str(value).lower()
                reader = reader.option(key, value)

            try:
                df = reader.load(full_path)

                if "filter" in options:
                    df = df.filter(options["filter"])
                    ctx.debug(f"Applied filter: {options['filter']}")

                elapsed = (time.time() - start_time) * 1000

                if not streaming:
                    partition_count = self._safe_partition_count(df)
                    ctx.log_spark_metrics(partition_count=partition_count)
                    ctx.log_file_io(path=path, format=format, mode="read")
                    ctx.info(
                        f"File read completed: {path}",
                        elapsed_ms=round(elapsed, 2),
                        partitions=partition_count,
                        format=format,
                    )
                else:
                    ctx.info(f"Streaming read started: {path}", elapsed_ms=round(elapsed, 2))

                return df

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    f"File read failed: {path}",
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                    format=format,
                )
                raise
        else:
            ctx.error("Either path or table must be provided")
            raise ValueError("Either path or table must be provided")

    def write(
        self,
        df: Any,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        register_table: Optional[str] = None,
        mode: str = "overwrite",
        options: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Write data using Spark.

        Args:
            df: Spark DataFrame to write
            connection: Connection object
            format: Output format (csv, parquet, json, delta)
            table: Table name
            path: File path
            register_table: Name to register as external table (if path is used)
            mode: Write mode (overwrite, append, error, ignore, upsert, append_once)
            options: Format-specific options (including partition_by for partitioning)
            streaming_config: StreamingWriteConfig for streaming DataFrames

        Returns:
            Optional dictionary containing Delta commit metadata (if format=delta),
            or streaming query info (if streaming)
        """
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()
        options = options or {}

        if getattr(df, "isStreaming", False) is True:
            return self._write_streaming(
                df=df,
                connection=connection,
                format=format,
                table=table,
                path=path,
                register_table=register_table,
                options=options,
                streaming_config=streaming_config,
            )

        target_identifier = table or path or "unknown"
        try:
            partition_count = df.rdd.getNumPartitions()
        except Exception:
            partition_count = 1  # Fallback for mocks or unsupported DataFrames

        # Auto-coalesce DataFrames for Delta writes to reduce file overhead
        # Use coalesce_partitions option to explicitly set target partitions
        # NOTE: We avoid df.count() here as it would trigger double-evaluation of lazy DataFrames
        coalesce_partitions = options.pop("coalesce_partitions", None)
        if (
            coalesce_partitions
            and isinstance(partition_count, int)
            and partition_count > coalesce_partitions
        ):
            df = df.coalesce(coalesce_partitions)
            ctx.debug(
                f"Coalesced DataFrame to {coalesce_partitions} partition(s)",
                original_partitions=partition_count,
            )
            partition_count = coalesce_partitions

        ctx.debug(
            "Starting Spark write",
            format=format,
            target=target_identifier,
            mode=mode,
            partitions=partition_count,
        )

        # SQL Server / Azure SQL Support
        if is_sql_format(format):
            if not hasattr(connection, "get_spark_options"):
                conn_type = type(connection).__name__
                msg = f"Connection type '{conn_type}' does not support Spark SQL write"
                ctx.error(msg, connection_type=conn_type)
                raise ValueError(msg)

            jdbc_options = connection.get_spark_options()
            merged_options = {**jdbc_options, **options}

            if table:
                merged_options["dbtable"] = table
            elif "dbtable" not in merged_options:
                ctx.error("SQL format requires 'table' config or 'dbtable' option")
                raise ValueError("SQL format requires 'table' config or 'dbtable' option")

            # Handle MERGE mode (SQL Server only)
            if mode == "merge":
                if getattr(connection, "sql_dialect", "mssql") != "mssql":
                    raise NotImplementedError(
                        f"SQL MERGE mode is currently only supported for SQL Server connections. "
                        f"Connection dialect '{getattr(connection, 'sql_dialect', 'unknown')}' "
                        f"does not support MERGE. Use mode='append' or mode='overwrite' instead."
                    )

                merge_keys = options.get("merge_keys")
                merge_options = options.get("merge_options")

                if not merge_keys:
                    ctx.error("MERGE mode requires 'merge_keys' in options")
                    raise ValueError(
                        "MERGE mode requires 'merge_keys' in options. "
                        "Specify the key columns for the MERGE ON clause."
                    )

                from odibi.writers.sql_server_writer import SqlServerMergeWriter

                writer = SqlServerMergeWriter(connection)

                # Check if bulk_copy is enabled for merge staging
                staging_connection = None
                if merge_options and getattr(merge_options, "bulk_copy", False):
                    staging_conn_name = getattr(merge_options, "staging_connection", None)
                    if not staging_conn_name:
                        raise ValueError(
                            "bulk_copy=True requires 'staging_connection' to be set. "
                            "Specify an ADLS/Blob connection for staging files."
                        )
                    # Get staging connection from options (passed from node executor)
                    staging_connection = options.get("_staging_connection")
                    if not staging_connection:
                        raise ValueError(
                            f"Staging connection '{staging_conn_name}' not found. "
                            "Ensure the connection is defined in your project config."
                        )

                ctx.debug(
                    "Executing SQL Server MERGE",
                    target=table,
                    merge_keys=merge_keys,
                    bulk_copy=staging_connection is not None,
                )

                try:
                    result = writer.merge(
                        df=df,
                        spark_engine=self,
                        target_table=table,
                        merge_keys=merge_keys,
                        options=merge_options,
                        jdbc_options=jdbc_options,
                        staging_connection=staging_connection,
                    )
                    elapsed = (time.time() - start_time) * 1000
                    ctx.log_file_io(path=target_identifier, format=format, mode="write")
                    ctx.info(
                        "SQL Server MERGE completed",
                        target=target_identifier,
                        mode=mode,
                        inserted=result.inserted,
                        updated=result.updated,
                        deleted=result.deleted,
                        elapsed_ms=round(elapsed, 2),
                    )
                    return {
                        "mode": "merge",
                        "inserted": result.inserted,
                        "updated": result.updated,
                        "deleted": result.deleted,
                        "total_affected": result.total_affected,
                    }

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "SQL Server MERGE failed",
                        target=target_identifier,
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise

            # Handle enhanced overwrite with strategies
            if mode == "overwrite" and options.get("overwrite_options"):
                from odibi.writers.sql_server_writer import SqlServerMergeWriter

                overwrite_options = options.get("overwrite_options")
                writer = SqlServerMergeWriter(connection)

                # Check if bulk_copy is enabled
                use_bulk_copy = getattr(overwrite_options, "bulk_copy", False)

                if use_bulk_copy:
                    # Bulk copy mode - use staging files + BULK INSERT
                    staging_conn_name = getattr(overwrite_options, "staging_connection", None)
                    external_data_source = getattr(overwrite_options, "external_data_source", None)

                    if not staging_conn_name:
                        raise ValueError(
                            "bulk_copy=True requires 'staging_connection' to be set. "
                            "Specify an ADLS/Blob connection for staging files."
                        )
                    if not external_data_source:
                        raise ValueError(
                            "bulk_copy=True requires 'external_data_source' to be set. "
                            "This is the SQL Server external data source name pointing "
                            "to your staging storage."
                        )

                    # Get staging connection from options (passed from node executor)
                    staging_connection = options.get("_staging_connection")
                    if not staging_connection:
                        raise ValueError(
                            f"Staging connection '{staging_conn_name}' not found. "
                            "Ensure the connection is defined in your project config."
                        )

                    ctx.debug(
                        "Executing SQL Server bulk copy overwrite",
                        target=table,
                        staging_connection=staging_conn_name,
                        external_data_source=external_data_source,
                    )

                    # Get bulk copy context for staging path organization
                    bulk_copy_context = options.get("_bulk_copy_context", {})

                    try:
                        result = writer.bulk_copy_spark(
                            df=df,
                            target_table=table,
                            staging_connection=staging_connection,
                            external_data_source=external_data_source,
                            options=overwrite_options,
                            bulk_copy_context=bulk_copy_context,
                        )
                        elapsed = (time.time() - start_time) * 1000
                        ctx.log_file_io(path=target_identifier, format=format, mode="write")
                        ctx.info(
                            "SQL Server bulk copy completed",
                            target=target_identifier,
                            strategy="bulk_copy",
                            rows_written=result.rows_written,
                            elapsed_ms=round(elapsed, 2),
                        )
                        return {
                            "mode": "overwrite",
                            "strategy": "bulk_copy",
                            "rows_written": result.rows_written,
                        }

                    except Exception as e:
                        elapsed = (time.time() - start_time) * 1000
                        ctx.error(
                            "SQL Server bulk copy failed",
                            target=target_identifier,
                            error_type=type(e).__name__,
                            error_message=str(e),
                            elapsed_ms=round(elapsed, 2),
                        )
                        raise

                # Standard JDBC-based overwrite
                ctx.debug(
                    "Executing SQL Server enhanced overwrite",
                    target=table,
                    strategy=(
                        overwrite_options.strategy.value
                        if hasattr(overwrite_options, "strategy")
                        else "truncate_insert"
                    ),
                )

                try:
                    result = writer.overwrite_spark(
                        df=df,
                        target_table=table,
                        options=overwrite_options,
                        jdbc_options=jdbc_options,
                    )
                    elapsed = (time.time() - start_time) * 1000
                    ctx.log_file_io(path=target_identifier, format=format, mode="write")
                    ctx.info(
                        "SQL Server enhanced overwrite completed",
                        target=target_identifier,
                        strategy=result.strategy,
                        rows_written=result.rows_written,
                        elapsed_ms=round(elapsed, 2),
                    )
                    return {
                        "mode": "overwrite",
                        "strategy": result.strategy,
                        "rows_written": result.rows_written,
                    }

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "SQL Server enhanced overwrite failed",
                        target=target_identifier,
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise

            if mode not in ["overwrite", "append", "ignore", "error"]:
                if mode == "fail":
                    mode = "error"
                else:
                    ctx.error(f"Write mode '{mode}' not supported for Spark SQL write")
                    raise ValueError(f"Write mode '{mode}' not supported for Spark SQL write")

            ctx.debug("Executing JDBC write", target=table or merged_options.get("dbtable"))

            try:
                df.write.format("jdbc").options(**merged_options).mode(mode).save()
                elapsed = (time.time() - start_time) * 1000
                ctx.log_file_io(path=target_identifier, format=format, mode="write")
                ctx.info(
                    "JDBC write completed",
                    target=target_identifier,
                    mode=mode,
                    elapsed_ms=round(elapsed, 2),
                )
                return None

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "JDBC write failed",
                    target=target_identifier,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        # Unity Catalog: qualify bare table/path names via the connection so
        # the configured catalog.schema is honoured.  Without this, Spark
        # resolves bare names against the session default (workspace.default),
        # silently ignoring the YAML schema parameter.
        from odibi.connections.unity_catalog import UnityCatalogConnection

        if isinstance(connection, UnityCatalogConnection):
            if not table and path:
                resolved = connection.get_path(path)
                if not resolved.startswith("/"):
                    table = resolved
                    path = None
            elif table and "." not in table:
                table = connection.get_path(table)

        # Handle Upsert/AppendOnce (Delta Only)
        if mode in ["upsert", "append_once"]:
            if format != "delta":
                ctx.error(f"Mode '{mode}' only supported for Delta format")
                raise NotImplementedError(
                    f"Mode '{mode}' only supported for Delta format in Spark engine."
                )

            keys = options.get("keys")
            if not keys:
                ctx.error(f"Mode '{mode}' requires 'keys' list in options")
                raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

            if isinstance(keys, str):
                keys = [keys]

            exists = self.table_exists(connection, table, path)
            ctx.debug("Table existence check for merge", target=target_identifier, exists=exists)

            if not exists:
                mode = "overwrite"
                ctx.debug("Target does not exist, falling back to overwrite mode")
            else:
                from delta.tables import DeltaTable

                target_dt = None
                target_name = ""
                is_table_target = False

                if table:
                    target_dt = DeltaTable.forName(self.spark, table)
                    target_name = table
                    is_table_target = True
                elif path:
                    full_path = connection.get_path(path)
                    target_dt = DeltaTable.forPath(self.spark, full_path)
                    target_name = full_path
                    is_table_target = False

                condition = " AND ".join([f"target.`{k}` = source.`{k}`" for k in keys])
                ctx.debug("Executing Delta merge", merge_mode=mode, keys=keys, condition=condition)

                merge_builder = target_dt.alias("target").merge(df.alias("source"), condition)

                try:
                    if mode == "upsert":
                        merge_builder.whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()
                    elif mode == "append_once":
                        merge_builder.whenNotMatchedInsertAll().execute()

                    elapsed = (time.time() - start_time) * 1000
                    ctx.info(
                        "Delta merge completed",
                        target=target_name,
                        mode=mode,
                        elapsed_ms=round(elapsed, 2),
                    )

                    self._optimize_delta_write(target_name, options, is_table=is_table_target)
                    commit_info = self._get_last_delta_commit_info(
                        target_name, is_table=is_table_target
                    )

                    if commit_info:
                        ctx.debug(
                            "Delta commit info",
                            version=commit_info.get("version"),
                            operation=commit_info.get("operation"),
                        )

                    return commit_info

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "Delta merge failed",
                        target=target_name,
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise

        # Get output location
        if table:
            # Managed/External Table (Catalog)
            ctx.debug(f"Writing to catalog table: {table}")
            writer = df.write.format(format).mode(mode)

            partition_by = options.get("partition_by")
            if partition_by:
                if isinstance(partition_by, str):
                    partition_by = [partition_by]
                writer = writer.partitionBy(*partition_by)
                ctx.debug(f"Partitioning by: {partition_by}")

            for key, value in options.items():
                writer = writer.option(key, value)

            try:
                writer.saveAsTable(table)
                elapsed = (time.time() - start_time) * 1000

                ctx.log_file_io(
                    path=table,
                    format=format,
                    mode=mode,
                    partitions=partition_by,
                )
                ctx.info(
                    f"Table write completed: {table}",
                    mode=mode,
                    elapsed_ms=round(elapsed, 2),
                )

                if format == "delta":
                    self._optimize_delta_write(table, options, is_table=True)
                    return self._get_last_delta_commit_info(table, is_table=True)
                return None

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    f"Table write failed: {table}",
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        elif path:
            full_path = connection.get_path(path)
        else:
            ctx.error("Either path or table must be provided")
            raise ValueError("Either path or table must be provided")

        # Extract partition_by option
        partition_by = options.pop("partition_by", None) or options.pop("partitionBy", None)

        # Extract cluster_by option (Liquid Clustering)
        cluster_by = options.pop("cluster_by", None)

        # Warn about partitioning anti-patterns
        if partition_by and cluster_by:
            import warnings

            ctx.warning(
                "Conflict: Both 'partition_by' and 'cluster_by' are set",
                partition_by=partition_by,
                cluster_by=cluster_by,
            )
            warnings.warn(
                "⚠️  Conflict: Both 'partition_by' and 'cluster_by' (Liquid Clustering) are set. "
                "Liquid Clustering supersedes partitioning. 'partition_by' will be ignored "
                "if the table is being created now.",
                UserWarning,
            )

        elif partition_by:
            import warnings

            ctx.warning(
                "Partitioning warning: ensure low-cardinality columns",
                partition_by=partition_by,
            )
            warnings.warn(
                "⚠️  Partitioning can cause performance issues if misused. "
                "Only partition on low-cardinality columns (< 1000 unique values) "
                "and ensure each partition has > 1000 rows.",
                UserWarning,
            )

        # Handle Upsert/Append-Once for Delta Lake (Path-based only for now)
        if format == "delta" and mode in ["upsert", "append_once"]:
            try:
                from delta.tables import DeltaTable
            except ImportError:
                ctx.error("Delta Lake support requires 'delta-spark'")
                raise ImportError("Delta Lake support requires 'delta-spark'")

            if "keys" not in options:
                ctx.error(f"Mode '{mode}' requires 'keys' list in options")
                raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

            if DeltaTable.isDeltaTable(self.spark, full_path):
                ctx.debug(f"Performing Delta merge at path: {full_path}")
                delta_table = DeltaTable.forPath(self.spark, full_path)
                keys = options["keys"]
                if isinstance(keys, str):
                    keys = [keys]

                condition = " AND ".join([f"target.{k} = source.{k}" for k in keys])
                merger = delta_table.alias("target").merge(df.alias("source"), condition)

                try:
                    if mode == "upsert":
                        merger.whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()
                    else:
                        merger.whenNotMatchedInsertAll().execute()

                    elapsed = (time.time() - start_time) * 1000
                    ctx.info(
                        "Delta merge completed at path",
                        path=path,
                        mode=mode,
                        elapsed_ms=round(elapsed, 2),
                    )

                    if register_table:
                        try:
                            table_in_catalog = self.spark.catalog.tableExists(register_table)
                            needs_registration = not table_in_catalog

                            # Handle orphan catalog entries (only for path-not-found errors)
                            if table_in_catalog:
                                try:
                                    # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
                                    self.spark.table(register_table).limit(1).collect()
                                    ctx.debug(
                                        f"Table '{register_table}' already registered and valid"
                                    )
                                except Exception as verify_err:
                                    error_str = str(verify_err)
                                    is_orphan = (
                                        "DELTA_PATH_DOES_NOT_EXIST" in error_str
                                        or "Path does not exist" in error_str
                                        or "FileNotFoundException" in error_str
                                    )
                                    if is_orphan:
                                        ctx.warning(
                                            f"Table '{register_table}' is orphan, re-registering"
                                        )
                                        try:
                                            self.spark.sql(f"DROP TABLE IF EXISTS {register_table}")
                                        except Exception:
                                            pass
                                        needs_registration = True
                                    else:
                                        ctx.debug(
                                            f"Table '{register_table}' verify failed, "
                                            "skipping registration"
                                        )

                            if needs_registration:
                                create_sql = (
                                    f"CREATE TABLE IF NOT EXISTS {register_table} "
                                    f"USING DELTA LOCATION '{full_path}'"
                                )
                                self.spark.sql(create_sql)
                                ctx.info(f"Registered table: {register_table}", path=full_path)
                        except Exception as e:
                            ctx.error(
                                f"Failed to register external table '{register_table}'",
                                error_message=str(e),
                            )

                    self._optimize_delta_write(full_path, options, is_table=False)
                    return self._get_last_delta_commit_info(full_path, is_table=False)

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "Delta merge failed at path",
                        path=path,
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise
            else:
                mode = "overwrite"
                ctx.debug("Target does not exist, falling back to overwrite mode")

        # Write based on format (Path-based)
        ctx.debug(f"Writing to path: {full_path}")

        # Handle Liquid Clustering (New Table Creation via SQL)
        if format == "delta" and cluster_by:
            should_create = False
            target_name = None

            if table:
                target_name = table
                if mode == "overwrite":
                    should_create = True
                elif mode == "append":
                    if not self.spark.catalog.tableExists(table):
                        should_create = True
            elif path:
                full_path = connection.get_path(path)
                target_name = f"delta.`{full_path}`"
                if mode == "overwrite":
                    should_create = True
                elif mode == "append":
                    try:
                        from delta.tables import DeltaTable

                        if not DeltaTable.isDeltaTable(self.spark, full_path):
                            should_create = True
                    except ImportError:
                        pass

            if should_create:
                if isinstance(cluster_by, str):
                    cluster_by = [cluster_by]

                cols = ", ".join(cluster_by)
                temp_view = f"odibi_temp_writer_{abs(hash(str(target_name)))}"
                df.createOrReplaceTempView(temp_view)

                create_cmd = (
                    "CREATE OR REPLACE TABLE"
                    if mode == "overwrite"
                    else "CREATE TABLE IF NOT EXISTS"
                )

                sql = (
                    f"{create_cmd} {target_name} USING DELTA CLUSTER BY ({cols}) "
                    f"AS SELECT * FROM {temp_view}"
                )

                ctx.debug("Creating clustered Delta table", sql=sql, cluster_by=cluster_by)

                try:
                    self.spark.sql(sql)
                    self.spark.catalog.dropTempView(temp_view)

                    elapsed = (time.time() - start_time) * 1000
                    ctx.info(
                        "Clustered Delta table created",
                        target=target_name,
                        cluster_by=cluster_by,
                        elapsed_ms=round(elapsed, 2),
                    )

                    if register_table and path:
                        try:
                            reg_sql = (
                                f"CREATE TABLE IF NOT EXISTS {register_table} "
                                f"USING DELTA LOCATION '{full_path}'"
                            )
                            self.spark.sql(reg_sql)
                            ctx.info(f"Registered table: {register_table}")
                        except Exception:
                            pass

                    if format == "delta":
                        self._optimize_delta_write(
                            target_name if table else full_path, options, is_table=bool(table)
                        )
                        return self._get_last_delta_commit_info(
                            target_name if table else full_path, is_table=bool(table)
                        )
                    return None

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "Failed to create clustered Delta table",
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise

        # Extract table_properties from options
        table_properties = options.pop("table_properties", None)

        # For column mapping and other properties that must be set BEFORE write
        original_configs = {}
        if table_properties and format == "delta":
            for prop_name, prop_value in table_properties.items():
                spark_conf_key = (
                    f"spark.databricks.delta.properties.defaults.{prop_name.replace('delta.', '')}"
                )
                try:
                    original_configs[spark_conf_key] = self.spark.conf.get(spark_conf_key, None)
                except Exception:
                    original_configs[spark_conf_key] = None
                self.spark.conf.set(spark_conf_key, prop_value)
            ctx.debug(
                "Applied table properties as session defaults",
                properties=list(table_properties.keys()),
            )

        writer = df.write.format(format).mode(mode)

        if partition_by:
            if isinstance(partition_by, str):
                partition_by = [partition_by]
            writer = writer.partitionBy(*partition_by)
            ctx.debug(f"Partitioning by: {partition_by}")

        for key, value in options.items():
            writer = writer.option(key, value)

        try:
            writer.save(full_path)
            elapsed = (time.time() - start_time) * 1000

            ctx.log_file_io(
                path=path,
                format=format,
                mode=mode,
                partitions=partition_by,
            )
            ctx.info(
                f"File write completed: {path}",
                format=format,
                mode=mode,
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                f"File write failed: {path}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise
        finally:
            for conf_key, original_value in original_configs.items():
                if original_value is None:
                    self.spark.conf.unset(conf_key)
                else:
                    self.spark.conf.set(conf_key, original_value)

        if format == "delta":
            self._optimize_delta_write(full_path, options, is_table=False)

        if register_table and format == "delta":
            try:
                table_in_catalog = self.spark.catalog.tableExists(register_table)
                needs_registration = not table_in_catalog

                # Handle orphan catalog entries: table exists but points to deleted path
                # Only treat as orphan if it's specifically a DELTA_PATH_DOES_NOT_EXIST error
                if table_in_catalog:
                    try:
                        # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
                        self.spark.table(register_table).limit(1).collect()
                        ctx.debug(
                            f"Table '{register_table}' already registered and valid, "
                            "skipping registration"
                        )
                    except Exception as verify_err:
                        error_str = str(verify_err)
                        is_orphan = (
                            "DELTA_PATH_DOES_NOT_EXIST" in error_str
                            or "Path does not exist" in error_str
                            or "FileNotFoundException" in error_str
                        )

                        if is_orphan:
                            # Orphan entry - table in catalog but path was deleted
                            ctx.warning(
                                f"Table '{register_table}' is orphan (path deleted), "
                                "dropping and re-registering",
                                error_message=error_str[:200],
                            )
                            try:
                                self.spark.sql(f"DROP TABLE IF EXISTS {register_table}")
                            except Exception:
                                pass  # Best effort cleanup
                            needs_registration = True
                        else:
                            # Other error (auth, network, etc.) - don't drop, just log
                            ctx.debug(
                                f"Table '{register_table}' exists but verify failed "
                                "(not orphan), skipping registration",
                                error_message=error_str[:200],
                            )

                if needs_registration:
                    ctx.debug(f"Registering table '{register_table}' at '{full_path}'")
                    reg_sql = (
                        f"CREATE TABLE IF NOT EXISTS {register_table} "
                        f"USING DELTA LOCATION '{full_path}'"
                    )
                    self.spark.sql(reg_sql)
                    ctx.info(f"Registered table: {register_table}", path=full_path)
            except Exception as e:
                ctx.error(
                    f"Failed to register table '{register_table}'",
                    error_message=str(e),
                )
                raise RuntimeError(
                    f"Failed to register external table '{register_table}': {e}"
                ) from e

        if format == "delta":
            return self._get_last_delta_commit_info(full_path, is_table=False)

        return None

    def _write_streaming(
        self,
        df,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        register_table: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Any] = None,
    ) -> Dict[str, Any]:
        """Write streaming DataFrame using Spark Structured Streaming.

        Args:
            df: Streaming Spark DataFrame
            connection: Connection object
            format: Output format (delta, kafka, etc.)
            table: Table name
            path: File path
            register_table: Name to register as external table (if path is used)
            options: Format-specific options
            streaming_config: StreamingWriteConfig with streaming parameters

        Returns:
            Dictionary with streaming query information
        """
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()
        options = options or {}

        if streaming_config is None:
            ctx.error("Streaming DataFrame requires streaming_config")
            raise ValueError(
                "Streaming DataFrame detected but no streaming_config provided. "
                "Add a 'streaming' section to your write config with at least "
                "'checkpoint_location' specified."
            )

        target_identifier = table or path or "unknown"

        checkpoint_location = streaming_config.checkpoint_location
        if checkpoint_location and connection:
            if not checkpoint_location.startswith(
                ("abfss://", "s3://", "gs://", "dbfs://", "hdfs://", "wasbs://")
            ):
                checkpoint_location = connection.get_path(checkpoint_location)
                ctx.debug(
                    "Resolved checkpoint location through connection",
                    original=streaming_config.checkpoint_location,
                    resolved=checkpoint_location,
                )

        # Extract table_properties from options for Delta column mapping etc.
        table_properties = options.pop("table_properties", None)

        # For column mapping and other properties that must be set BEFORE write
        original_configs = {}
        if table_properties and format == "delta":
            for prop_name, prop_value in table_properties.items():
                spark_conf_key = (
                    f"spark.databricks.delta.properties.defaults.{prop_name.replace('delta.', '')}"
                )
                try:
                    original_configs[spark_conf_key] = self.spark.conf.get(spark_conf_key, None)
                except Exception:
                    original_configs[spark_conf_key] = None
                self.spark.conf.set(spark_conf_key, prop_value)
            ctx.debug(
                "Applied table properties as session defaults for streaming",
                properties=list(table_properties.keys()),
            )

        ctx.debug(
            "Starting streaming write",
            format=format,
            target=target_identifier,
            output_mode=streaming_config.output_mode,
            checkpoint=checkpoint_location,
        )

        writer = df.writeStream.format(format)
        writer = writer.outputMode(streaming_config.output_mode)
        writer = writer.option("checkpointLocation", checkpoint_location)

        if streaming_config.query_name:
            writer = writer.queryName(streaming_config.query_name)

        if streaming_config.trigger:
            trigger = streaming_config.trigger
            if trigger.once:
                writer = writer.trigger(once=True)
            elif trigger.available_now:
                writer = writer.trigger(availableNow=True)
            elif trigger.processing_time:
                writer = writer.trigger(processingTime=trigger.processing_time)
            elif trigger.continuous:
                writer = writer.trigger(continuous=trigger.continuous)

        partition_by = options.pop("partition_by", None) or options.pop("partitionBy", None)
        if partition_by:
            if isinstance(partition_by, str):
                partition_by = [partition_by]
            writer = writer.partitionBy(*partition_by)
            ctx.debug(f"Partitioning by: {partition_by}")

        for key, value in options.items():
            writer = writer.option(key, value)

        # Capture Delta version before streaming to detect if new data was written
        version_before = None
        if path and format == "delta":
            try:
                full_path = connection.get_path(path)
                history_df = self.spark.sql(f"DESCRIBE HISTORY delta.`{full_path}` LIMIT 1")
                rows = history_df.collect()
                if rows:
                    version_before = (
                        rows[0].version if hasattr(rows[0], "version") else rows[0]["version"]
                    )
            except Exception:
                pass  # Table may not exist yet

        try:
            if table:
                query = writer.toTable(table)
                ctx.info(
                    f"Streaming query started: writing to table {table}",
                    query_id=str(query.id),
                    query_name=query.name,
                )
            elif path:
                full_path = connection.get_path(path)
                query = writer.start(full_path)
                ctx.info(
                    f"Streaming query started: writing to path {path}",
                    query_id=str(query.id),
                    query_name=query.name,
                )
            else:
                ctx.error("Either path or table must be provided for streaming write")
                raise ValueError(
                    "Streaming write operation failed: neither 'path' nor 'table' was provided. "
                    "Specify a file path or table name in your streaming configuration."
                )

            elapsed = (time.time() - start_time) * 1000

            result = {
                "streaming": True,
                "query_id": str(query.id),
                "query_name": query.name,
                "status": "running",
                "target": target_identifier,
                "output_mode": streaming_config.output_mode,
                "checkpoint_location": streaming_config.checkpoint_location,
                "elapsed_ms": round(elapsed, 2),
            }

            should_wait = streaming_config.await_termination
            if streaming_config.trigger:
                trigger = streaming_config.trigger
                if trigger.once or trigger.available_now:
                    should_wait = True

            if should_wait:
                ctx.info(
                    "Awaiting streaming query termination",
                    timeout_seconds=streaming_config.timeout_seconds,
                )
                query.awaitTermination(streaming_config.timeout_seconds)
                result["status"] = "terminated"
                elapsed = (time.time() - start_time) * 1000
                result["elapsed_ms"] = round(elapsed, 2)

                # Get rows written from streaming query progress
                rows_written = None

                # Method 1: Sum all micro-batch progress (most accurate for streaming)
                # For autoloader, numInputRows = rows read from source files in this run
                try:
                    recent_progress = query.recentProgress
                    if recent_progress:
                        total_rows = 0
                        for prog in recent_progress:
                            if prog:
                                # Use numInputRows - this is rows processed from source
                                # NOT sink.numOutputRows which can be cumulative
                                input_rows = prog.get("numInputRows")
                                if input_rows is not None:
                                    total_rows += input_rows
                        if total_rows > 0:
                            rows_written = total_rows
                except Exception:
                    pass

                # Method 1b: Try lastProgress if recentProgress is empty
                if rows_written is None:
                    try:
                        last_progress = query.lastProgress
                        if last_progress:
                            input_rows = last_progress.get("numInputRows")
                            if input_rows is not None and input_rows > 0:
                                rows_written = input_rows
                    except Exception:
                        pass

                # Method 2: Check Delta table history for rows written in THIS run only
                # Compare version before/after to detect if new data was written
                if rows_written is None and path and format == "delta":
                    try:
                        full_path = connection.get_path(path)
                        history_df = self.spark.sql(f"DESCRIBE HISTORY delta.`{full_path}` LIMIT 1")
                        rows = history_df.collect()
                        if rows:
                            row = rows[0]
                            version_after = (
                                row.version if hasattr(row, "version") else row["version"]
                            )

                            # If version hasn't changed, no new data was written
                            if version_before is not None and version_after == version_before:
                                rows_written = 0
                            else:
                                # New version created - get the row count from metrics
                                metrics = (
                                    row.operationMetrics
                                    if hasattr(row, "operationMetrics")
                                    else row["operationMetrics"]
                                )
                                if metrics:
                                    # For streaming: numAddedRows is rows added in this batch
                                    added_rows = metrics.get("numAddedRows")
                                    if added_rows is not None:
                                        rows_written = int(added_rows)
                                    else:
                                        # Fallback to numOutputRows
                                        output_rows = metrics.get("numOutputRows")
                                        if output_rows is not None:
                                            rows_written = int(output_rows)
                    except Exception:
                        pass

                if rows_written is not None:
                    result["_cached_row_count"] = rows_written

                ctx.info(
                    "Streaming query terminated",
                    query_id=str(query.id),
                    elapsed_ms=round(elapsed, 2),
                    rows_written=rows_written,
                )

                if register_table and path and format == "delta":
                    full_path = connection.get_path(path)
                    try:
                        self.spark.sql(
                            f"CREATE TABLE IF NOT EXISTS {register_table} "
                            f"USING DELTA LOCATION '{full_path}'"
                        )
                        ctx.info(
                            f"Registered external table: {register_table}",
                            path=full_path,
                        )
                        result["registered_table"] = register_table
                    except Exception as reg_err:
                        ctx.warning(
                            f"Failed to register external table '{register_table}'",
                            error=str(reg_err),
                        )
            else:
                result["streaming_query"] = query
                if register_table:
                    ctx.warning(
                        "register_table ignored for continuous streaming. "
                        "Table will be registered after query terminates or manually."
                    )

            return result

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "Streaming write failed",
                target=target_identifier,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise
        finally:
            # Restore original session configs
            for conf_key, original_value in original_configs.items():
                try:
                    if original_value is None:
                        self.spark.conf.unset(conf_key)
                    else:
                        self.spark.conf.set(conf_key, original_value)
                except Exception:
                    pass  # Best effort cleanup

    def execute_sql(self, sql: str, context: Any = None) -> Any:
        """Execute SQL query in Spark.

        Args:
            sql: SQL query string
            context: Execution context (optional, not used for Spark)

        Returns:
            Spark DataFrame with query results
        """
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        ctx.debug("Executing Spark SQL", query_preview=sql[:200] if len(sql) > 200 else sql)

        try:
            result = self.spark.sql(sql)
            elapsed = (time.time() - start_time) * 1000
            partition_count = self._safe_partition_count(result)

            ctx.log_spark_metrics(partition_count=partition_count)
            ctx.info(
                "Spark SQL executed",
                elapsed_ms=round(elapsed, 2),
                partitions=partition_count,
            )

            return result

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            error_type = type(e).__name__
            clean_message = _extract_spark_error_message(e)

            if "AnalysisException" in error_type:
                ctx.error(
                    "Spark SQL Analysis Error",
                    error_type=error_type,
                    error_message=clean_message,
                    query_preview=sql[:200] if len(sql) > 200 else sql,
                    elapsed_ms=round(elapsed, 2),
                )
                raise TransformError(f"Spark SQL Analysis Error: {clean_message}")

            if "ParseException" in error_type:
                ctx.error(
                    "Spark SQL Parse Error",
                    error_type=error_type,
                    error_message=clean_message,
                    query_preview=sql[:200] if len(sql) > 200 else sql,
                    elapsed_ms=round(elapsed, 2),
                )
                raise TransformError(f"Spark SQL Parse Error: {clean_message}")

            ctx.error(
                "Spark SQL execution failed",
                error_type=error_type,
                error_message=clean_message,
                elapsed_ms=round(elapsed, 2),
            )
            raise TransformError(f"Spark SQL Error: {clean_message}")

    def execute_transform(self, *args, **kwargs):
        raise NotImplementedError(
            "SparkEngine.execute_transform() will be implemented in Phase 2B. "
            "See PHASES.md for implementation plan."
        )

    def execute_operation(self, operation: str, params: Dict[str, Any], df) -> Any:
        """Execute built-in operation on Spark DataFrame."""
        ctx = get_logging_context().with_context(engine="spark")
        params = params or {}

        ctx.debug(f"Executing operation: {operation}", params=list(params.keys()))

        if operation == "pivot":
            group_by = params.get("group_by", [])
            pivot_column = params.get("pivot_column")
            value_column = params.get("value_column")
            agg_func = params.get("agg_func", "first")

            if not pivot_column or not value_column:
                ctx.error("Pivot requires 'pivot_column' and 'value_column'")
                raise ValueError("Pivot requires 'pivot_column' and 'value_column'")

            if isinstance(group_by, str):
                group_by = [group_by]

            agg_expr = {value_column: agg_func}
            return df.groupBy(*group_by).pivot(pivot_column).agg(agg_expr)

        elif operation == "drop_duplicates":
            subset = params.get("subset")
            if subset:
                if isinstance(subset, str):
                    subset = [subset]
                return df.dropDuplicates(subset=subset)
            return df.dropDuplicates()

        elif operation == "fillna":
            value = params.get("value")
            subset = params.get("subset")
            return df.fillna(value, subset=subset)

        elif operation == "drop":
            columns = params.get("columns")
            if not columns:
                return df
            if isinstance(columns, str):
                columns = [columns]
            return df.drop(*columns)

        elif operation == "rename":
            columns = params.get("columns")
            if not columns:
                return df

            res = df
            for old_name, new_name in columns.items():
                res = res.withColumnRenamed(old_name, new_name)
            return res

        elif operation == "sort":
            by = params.get("by")
            ascending = params.get("ascending", True)

            if not by:
                return df

            if isinstance(by, str):
                by = [by]

            if not ascending:
                from pyspark.sql.functions import desc

                sort_cols = [desc(c) for c in by]
                return df.orderBy(*sort_cols)

            return df.orderBy(*by)

        elif operation == "sample":
            fraction = params.get("frac", 0.1)
            seed = params.get("random_state")
            with_replacement = params.get("replace", False)
            return df.sample(withReplacement=with_replacement, fraction=fraction, seed=seed)

        else:
            # Fallback: check if operation is a registered transformer
            from odibi.context import EngineContext
            from odibi.registry import FunctionRegistry

            ctx.debug(
                f"Checking registry for operation: {operation}",
                registered_functions=list(FunctionRegistry._functions.keys())[:10],
                has_function=FunctionRegistry.has_function(operation),
            )

            if FunctionRegistry.has_function(operation):
                ctx.debug(f"Executing registered transformer as operation: {operation}")
                func = FunctionRegistry.get_function(operation)
                param_model = FunctionRegistry.get_param_model(operation)

                # Create EngineContext from current df
                from odibi.context import SparkContext

                engine_ctx = EngineContext(
                    context=SparkContext(self.spark),
                    df=df,
                    engine=self,
                    engine_type=self.engine_type,
                )

                # Validate and instantiate params
                if param_model:
                    validated_params = param_model(**params)
                    result_ctx = func(engine_ctx, validated_params)
                else:
                    result_ctx = func(engine_ctx, **params)

                return result_ctx.df

            ctx.error(f"Unsupported operation for Spark engine: {operation}")
            raise ValueError(f"Unsupported operation for Spark engine: {operation}")

    def count_nulls(self, df, columns: List[str]) -> Dict[str, int]:
        """Count nulls in specified columns."""
        from pyspark.sql.functions import col, count, when

        missing = set(columns) - set(df.columns)
        if missing:
            raise ValueError(f"Columns not found in DataFrame: {', '.join(missing)}")

        # Escape column names with backticks to handle special characters (., ,, spaces)
        aggs = [count(when(col(f"`{c}`").isNull(), c)).alias(c) for c in columns]
        rows = df.select(*aggs).collect()
        return rows[0].asDict() if rows else {c: 0 for c in columns}

    def validate_schema(self, df, schema_rules: Dict[str, Any]) -> List[str]:
        """Validate DataFrame schema."""
        failures = []

        if "required_columns" in schema_rules:
            required = schema_rules["required_columns"]
            missing = set(required) - set(df.columns)
            if missing:
                failures.append(f"Missing required columns: {', '.join(missing)}")

        if "types" in schema_rules:
            type_map = {
                "int": ["integer", "long", "short", "byte", "bigint"],
                "float": ["double", "float"],
                "str": ["string"],
                "bool": ["boolean"],
            }

            for col_name, expected_type in schema_rules["types"].items():
                if col_name not in df.columns:
                    failures.append(f"Column '{col_name}' not found for type validation")
                    continue

                actual_type = dict(df.dtypes)[col_name]
                expected_dtypes = type_map.get(expected_type, [expected_type])

                if actual_type not in expected_dtypes:
                    failures.append(
                        f"Column '{col_name}' has type '{actual_type}', expected '{expected_type}'"
                    )

        return failures

    def validate_data(self, df, validation_config: Any) -> List[str]:
        """Validate DataFrame against rules."""
        from pyspark.sql.functions import col

        ctx = get_logging_context().with_context(engine="spark")
        failures = []

        if validation_config.not_empty:
            if df.isEmpty():
                failures.append("DataFrame is empty")

        if validation_config.no_nulls:
            null_counts = self.count_nulls(df, validation_config.no_nulls)
            for col_name, count in null_counts.items():
                if count > 0:
                    failures.append(f"Column '{col_name}' has {count} null values")

        if validation_config.schema_validation:
            schema_failures = self.validate_schema(df, validation_config.schema_validation)
            failures.extend(schema_failures)

        if validation_config.ranges:
            for col_name, bounds in validation_config.ranges.items():
                if col_name in df.columns:
                    min_val = bounds.get("min")
                    max_val = bounds.get("max")

                    # Escape column names with backticks to handle special characters
                    escaped_col = f"`{col_name}`"
                    if min_val is not None:
                        count = df.filter(col(escaped_col) < min_val).count()
                        if count > 0:
                            failures.append(f"Column '{col_name}' has values < {min_val}")

                    if max_val is not None:
                        count = df.filter(col(escaped_col) > max_val).count()
                        if count > 0:
                            failures.append(f"Column '{col_name}' has values > {max_val}")
                else:
                    failures.append(f"Column '{col_name}' not found for range validation")

        if validation_config.allowed_values:
            for col_name, allowed in validation_config.allowed_values.items():
                if col_name in df.columns:
                    # Escape column names with backticks to handle special characters
                    escaped_col = f"`{col_name}`"
                    count = df.filter(~col(escaped_col).isin(allowed)).count()
                    if count > 0:
                        failures.append(f"Column '{col_name}' has invalid values")
                else:
                    failures.append(f"Column '{col_name}' not found for allowed values validation")

        ctx.log_validation_result(
            passed=len(failures) == 0,
            rule_name="data_validation",
            failures=failures if failures else None,
        )

        return failures

    def get_sample(self, df, n: int = 10) -> List[Dict[str, Any]]:
        """Get sample rows as list of dictionaries."""
        return [row.asDict() for row in df.limit(n).collect()]

    def table_exists(
        self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
    ) -> bool:
        """Check if table or location exists.

        Handles orphan catalog entries where the table is registered but
        the underlying Delta path no longer exists.
        """
        ctx = get_logging_context().with_context(engine="spark")

        if table:
            try:
                if not self.spark.catalog.tableExists(table):
                    ctx.debug(f"Table does not exist: {table}")
                    return False
                # Table exists in catalog - verify it's actually readable
                # This catches orphan entries where path was deleted
                # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
                self.spark.table(table).limit(1).collect()
                ctx.debug(f"Table existence check: {table}", exists=True)
                return True
            except Exception as e:
                # Table exists in catalog but underlying data is gone (orphan entry)
                # This is expected during first-run detection - log at debug level
                ctx.debug(
                    f"Table {table} exists in catalog but is not accessible (treating as first run)",
                    error_message=str(e),
                )
                return False
        elif path:
            try:
                from delta.tables import DeltaTable

                full_path = connection.get_path(path)
                exists = DeltaTable.isDeltaTable(self.spark, full_path)
                ctx.debug(f"Delta table existence check: {path}", exists=exists)
                return exists
            except ImportError:
                try:
                    full_path = connection.get_path(path)
                    exists = (
                        self.spark.sparkContext._gateway.jvm.org.apache.hadoop.fs.FileSystem.get(
                            self.spark.sparkContext._jsc.hadoopConfiguration()
                        ).exists(
                            self.spark.sparkContext._gateway.jvm.org.apache.hadoop.fs.Path(
                                full_path
                            )
                        )
                    )
                    ctx.debug(f"Path existence check: {path}", exists=exists)
                    return exists
                except Exception as e:
                    ctx.warning(f"Path existence check failed: {path}", error_message=str(e))
                    return False
            except Exception as e:
                ctx.warning(f"Table existence check failed: {path}", error_message=str(e))
                return False
        return False

    def get_table_schema(
        self,
        connection: Any,
        table: Optional[str] = None,
        path: Optional[str] = None,
        format: Optional[str] = None,
    ) -> Optional[Dict[str, str]]:
        """Get schema of an existing table/file."""
        ctx = get_logging_context().with_context(engine="spark")

        try:
            if table:
                if self.spark.catalog.tableExists(table):
                    schema = self.get_schema(self.spark.table(table))
                    ctx.debug(f"Retrieved schema for table: {table}", columns=len(schema))
                    return schema
            elif path:
                full_path = connection.get_path(path)
                if format == "delta":
                    from delta.tables import DeltaTable

                    if DeltaTable.isDeltaTable(self.spark, full_path):
                        schema = self.get_schema(DeltaTable.forPath(self.spark, full_path).toDF())
                        ctx.debug(f"Retrieved Delta schema: {path}", columns=len(schema))
                        return schema
                elif format == "parquet":
                    schema = self.get_schema(self.spark.read.parquet(full_path))
                    ctx.debug(f"Retrieved Parquet schema: {path}", columns=len(schema))
                    return schema
                elif format:
                    schema = self.get_schema(self.spark.read.format(format).load(full_path))
                    ctx.debug(f"Retrieved schema: {path}", format=format, columns=len(schema))
                    return schema
        except Exception as e:
            ctx.warning(
                "Failed to get schema",
                table=table,
                path=path,
                error_message=str(e),
            )
        return None

    def vacuum_delta(
        self,
        connection: Any,
        path: str,
        retention_hours: int = 168,
    ) -> None:
        """VACUUM a Delta table to remove old files."""
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        ctx.debug(
            "Starting Delta VACUUM",
            path=path,
            retention_hours=retention_hours,
        )

        try:
            from delta.tables import DeltaTable
        except ImportError:
            ctx.error("Delta Lake support requires 'delta-spark'")
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[spark]' "
                "with delta-spark. "
                "See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        try:
            delta_table = DeltaTable.forPath(self.spark, full_path)
            delta_table.vacuum(retention_hours / 24.0)

            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                "Delta VACUUM completed",
                path=path,
                retention_hours=retention_hours,
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "Delta VACUUM failed",
                path=path,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    def get_delta_history(
        self, connection: Any, path: str, limit: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Get Delta table history."""
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        ctx.debug("Fetching Delta history", path=path, limit=limit)

        try:
            from delta.tables import DeltaTable
        except ImportError:
            ctx.error("Delta Lake support requires 'delta-spark'")
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[spark]' "
                "with delta-spark. "
                "See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        try:
            delta_table = DeltaTable.forPath(self.spark, full_path)
            history_df = delta_table.history(limit) if limit else delta_table.history()
            history = [row.asDict() for row in history_df.collect()]

            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                "Delta history retrieved",
                path=path,
                versions_returned=len(history),
                elapsed_ms=round(elapsed, 2),
            )

            return history

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "Failed to get Delta history",
                path=path,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    def restore_delta(self, connection: Any, path: str, version: int) -> None:
        """Restore Delta table to a specific version."""
        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        ctx.debug("Restoring Delta table", path=path, version=version)

        try:
            from delta.tables import DeltaTable
        except ImportError:
            ctx.error("Delta Lake support requires 'delta-spark'")
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[spark]' "
                "with delta-spark. "
                "See README.md for installation instructions."
            )

        full_path = connection.get_path(path)

        try:
            delta_table = DeltaTable.forPath(self.spark, full_path)
            delta_table.restoreToVersion(version)

            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                "Delta table restored",
                path=path,
                version=version,
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "Delta restore failed",
                path=path,
                version=version,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    def maintain_table(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        config: Optional[Any] = None,
    ) -> None:
        """Run table maintenance operations (optimize, vacuum)."""
        if format != "delta" or not config or not config.enabled:
            return

        ctx = get_logging_context().with_context(engine="spark")
        start_time = time.time()

        if table:
            target = table
        elif path:
            full_path = connection.get_path(path)
            target = f"delta.`{full_path}`"
        else:
            return

        ctx.debug("Starting table maintenance", target=target)

        try:
            ctx.debug(f"Running OPTIMIZE on {target}")
            self.spark.sql(f"OPTIMIZE {target}")

            retention = config.vacuum_retention_hours
            if retention is not None and retention > 0:
                ctx.debug(f"Running VACUUM on {target}", retention_hours=retention)
                self.spark.sql(f"VACUUM {target} RETAIN {retention} HOURS")

            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                "Table maintenance completed",
                target=target,
                vacuum_retention_hours=retention,
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.warning(
                f"Auto-optimize failed for {target}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )

    def get_source_files(self, df) -> List[str]:
        """Get list of source files that generated this DataFrame."""
        try:
            return df.inputFiles()
        except Exception:
            return []

    def profile_nulls(self, df) -> Dict[str, float]:
        """Calculate null percentage for each column."""
        from pyspark.sql.functions import col, mean, when

        aggs = []
        for c in df.columns:
            # Escape column names with backticks to handle special characters (., ,, spaces)
            escaped_col = f"`{c}`"
            aggs.append(mean(when(col(escaped_col).isNull(), 1).otherwise(0)).alias(c))

        if not aggs:
            return {}

        try:
            rows = df.select(*aggs).collect()
            return rows[0].asDict() if rows else {}
        except Exception:
            return {}

    def filter_greater_than(self, df, column: str, value: Any) -> Any:
        """Filter DataFrame where column > value.

        Automatically casts string columns to timestamp for proper comparison.
        Tries multiple date formats including Oracle-style (DD-MON-YY).
        """
        from pyspark.sql import functions as F
        from pyspark.sql.types import StringType

        # Escape column names with backticks to handle special characters
        escaped_col = f"`{column}`"
        col_type = df.schema[column].dataType
        if isinstance(col_type, StringType):
            ts_col = self._parse_string_to_timestamp(F.col(escaped_col))
            return df.filter(ts_col > value)
        return df.filter(F.col(escaped_col) > value)

    def _parse_string_to_timestamp(self, col):
        """Parse string column to timestamp, trying multiple formats.

        Supports:
        - ISO format: 2024-04-20 07:11:01
        - Oracle format: 20-APR-24 07:11:01.0 (handles uppercase months)
        """
        from pyspark.sql import functions as F

        result = F.to_timestamp(col)

        result = F.coalesce(result, F.to_timestamp(col, "yyyy-MM-dd HH:mm:ss"))
        result = F.coalesce(result, F.to_timestamp(col, "yyyy-MM-dd'T'HH:mm:ss"))
        result = F.coalesce(result, F.to_timestamp(col, "MM/dd/yyyy HH:mm:ss"))

        col_oracle = F.concat(
            F.substring(col, 1, 3),
            F.upper(F.substring(col, 4, 1)),
            F.lower(F.substring(col, 5, 2)),
            F.substring(col, 7, 100),
        )
        result = F.coalesce(result, F.to_timestamp(col_oracle, "dd-MMM-yy HH:mm:ss.S"))
        result = F.coalesce(result, F.to_timestamp(col_oracle, "dd-MMM-yy HH:mm:ss"))

        return result

    def filter_coalesce(self, df, col1: str, col2: str, op: str, value: Any) -> Any:
        """Filter using COALESCE(col1, col2) op value.

        Automatically casts string columns to timestamp for proper comparison.
        Tries multiple date formats including Oracle-style (DD-MON-YY).
        """
        from pyspark.sql import functions as F
        from pyspark.sql.types import StringType

        # Escape column names with backticks to handle special characters
        escaped_col1 = f"`{col1}`"
        escaped_col2 = f"`{col2}`"
        col1_type = df.schema[col1].dataType
        col2_type = df.schema[col2].dataType

        if isinstance(col1_type, StringType):
            c1 = self._parse_string_to_timestamp(F.col(escaped_col1))
        else:
            c1 = F.col(escaped_col1)

        if isinstance(col2_type, StringType):
            c2 = self._parse_string_to_timestamp(F.col(escaped_col2))
        else:
            c2 = F.col(escaped_col2)

        coalesced = F.coalesce(c1, c2)

        if op == ">":
            return df.filter(coalesced > value)
        elif op == ">=":
            return df.filter(coalesced >= value)
        elif op == "<":
            return df.filter(coalesced < value)
        elif op == "<=":
            return df.filter(coalesced <= value)
        elif op == "=":
            return df.filter(coalesced == value)
        else:
            return df.filter(f"COALESCE({col1}, {col2}) {op} '{value}'")

    def add_write_metadata(
        self,
        df: Any,
        metadata_config: Any,
        source_connection: Optional[str] = None,
        source_table: Optional[str] = None,
        source_path: Optional[str] = None,
        is_file_source: bool = False,
    ) -> Any:
        """Add metadata columns to DataFrame before writing (Bronze layer lineage).

        Args:
            df: Spark DataFrame
            metadata_config: WriteMetadataConfig or True (for all defaults)
            source_connection: Name of the source connection
            source_table: Name of the source table (SQL sources)
            source_path: Path of the source file (file sources)
            is_file_source: True if source is a file-based read

        Returns:
            DataFrame with metadata columns added
        """
        from pyspark.sql import functions as F

        from odibi.config import WriteMetadataConfig

        if metadata_config is True:
            config = WriteMetadataConfig()
        elif isinstance(metadata_config, WriteMetadataConfig):
            config = metadata_config
        else:
            return df

        if config.extracted_at:
            df = df.withColumn("_extracted_at", F.current_timestamp())

        if config.source_file and is_file_source and source_path:
            df = df.withColumn("_source_file", F.lit(source_path))

        if config.source_connection and source_connection:
            df = df.withColumn("_source_connection", F.lit(source_connection))

        if config.source_table and source_table:
            df = df.withColumn("_source_table", F.lit(source_table))

        return df

__init__(connections=None, spark_session=None, config=None)

Initialize Spark engine with import guard.

Parameters:

Name Type Description Default
connections Optional[Dict[str, Any]]

Dictionary of connection objects (for multi-account config)

None
spark_session Any

Existing SparkSession (optional, creates new if None)

None
config Optional[Dict[str, Any]]

Engine configuration (optional)

None

Raises:

Type Description
ImportError

If pyspark not installed

Source code in odibi/engine/spark_engine.py
def __init__(
    self,
    connections: Optional[Dict[str, Any]] = None,
    spark_session: Any = None,
    config: Optional[Dict[str, Any]] = None,
):
    """Initialize Spark engine with import guard.

    Args:
        connections: Dictionary of connection objects (for multi-account config)
        spark_session: Existing SparkSession (optional, creates new if None)
        config: Engine configuration (optional)

    Raises:
        ImportError: If pyspark not installed
    """
    ctx = get_logging_context().with_context(engine="spark")
    ctx.debug("Initializing SparkEngine", connections_count=len(connections or {}))

    try:
        from pyspark.sql import SparkSession
    except ImportError as e:
        ctx.error(
            "PySpark not installed",
            error_type="ImportError",
            suggestion="pip install odibi[spark]",
        )
        raise ImportError(
            "Spark support requires 'pip install odibi[spark]'. "
            "See docs/setup_databricks.md for setup instructions."
        ) from e

    start_time = time.time()

    _reused_active = False
    if spark_session:
        self.spark = spark_session
    elif SparkSession.getActiveSession() is not None:
        # Reuse the SparkSession already running in this process. Every
        # notebook context has one — Databricks serverless and classic
        # clusters, Jupyter+Spark, an existing Spark Connect client. Building
        # a fresh session here instead would spin up a local Spark or a new
        # Connect client with no sc:// URL, which is the INVALID_CONNECT_URL
        # failure on serverless (and configure_spark_with_delta_pip assumes
        # local Spark). Only build below when there is genuinely no session.
        self.spark = SparkSession.getActiveSession()
        _reused_active = True
        ctx.debug("Reusing active SparkSession (no new session created)")
    else:
        # Configure Delta Lake support
        try:
            from delta import configure_spark_with_delta_pip

            builder = SparkSession.builder.appName("odibi").config(
                "spark.sql.sources.partitionOverwriteMode", "dynamic"
            )

            # Performance Optimizations
            builder = builder.config("spark.sql.execution.arrow.pyspark.enabled", "true")
            builder = builder.config("spark.sql.adaptive.enabled", "true")

            # Reduce Verbosity
            builder = builder.config(
                "spark.driver.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
            )
            builder = builder.config(
                "spark.executor.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
            )

            self.spark = configure_spark_with_delta_pip(builder).getOrCreate()
            try:
                self.spark.sparkContext.setLogLevel("ERROR")
            except Exception:
                pass  # Spark Connect: sparkContext not available

            ctx.debug("Delta Lake support enabled")

        except ImportError:
            ctx.debug("Delta Lake not available, using standard Spark")
            builder = SparkSession.builder.appName("odibi").config(
                "spark.sql.sources.partitionOverwriteMode", "dynamic"
            )

            # Performance Optimizations
            builder = builder.config("spark.sql.execution.arrow.pyspark.enabled", "true")
            builder = builder.config("spark.sql.adaptive.enabled", "true")

            # Reduce Verbosity
            builder = builder.config(
                "spark.driver.extraJavaOptions", "-Dlog4j.rootCategory=ERROR"
            )

            self.spark = builder.getOrCreate()
            try:
                self.spark.sparkContext.setLogLevel("ERROR")
            except Exception:
                pass  # Spark Connect: sparkContext not available

    self.config = config or {}
    self.connections = connections or {}

    # Configure all ADLS connections upfront
    self._configure_all_connections()

    # Apply user-defined Spark configs from performance settings
    self._apply_spark_config()

    elapsed = (time.time() - start_time) * 1000
    try:
        app_name = self.spark.sparkContext.appName
    except Exception:
        app_name = "unknown"  # Spark Connect: sparkContext not available
    ctx.info(
        "SparkEngine initialized",
        elapsed_ms=round(elapsed, 2),
        app_name=app_name,
        spark_version=self.spark.version,
        connections_configured=len(self.connections),
        using_existing_session=spark_session is not None or _reused_active,
    )

add_write_metadata(df, metadata_config, source_connection=None, source_table=None, source_path=None, is_file_source=False)

Add metadata columns to DataFrame before writing (Bronze layer lineage).

Parameters:

Name Type Description Default
df Any

Spark DataFrame

required
metadata_config Any

WriteMetadataConfig or True (for all defaults)

required
source_connection Optional[str]

Name of the source connection

None
source_table Optional[str]

Name of the source table (SQL sources)

None
source_path Optional[str]

Path of the source file (file sources)

None
is_file_source bool

True if source is a file-based read

False

Returns:

Type Description
Any

DataFrame with metadata columns added

Source code in odibi/engine/spark_engine.py
def add_write_metadata(
    self,
    df: Any,
    metadata_config: Any,
    source_connection: Optional[str] = None,
    source_table: Optional[str] = None,
    source_path: Optional[str] = None,
    is_file_source: bool = False,
) -> Any:
    """Add metadata columns to DataFrame before writing (Bronze layer lineage).

    Args:
        df: Spark DataFrame
        metadata_config: WriteMetadataConfig or True (for all defaults)
        source_connection: Name of the source connection
        source_table: Name of the source table (SQL sources)
        source_path: Path of the source file (file sources)
        is_file_source: True if source is a file-based read

    Returns:
        DataFrame with metadata columns added
    """
    from pyspark.sql import functions as F

    from odibi.config import WriteMetadataConfig

    if metadata_config is True:
        config = WriteMetadataConfig()
    elif isinstance(metadata_config, WriteMetadataConfig):
        config = metadata_config
    else:
        return df

    if config.extracted_at:
        df = df.withColumn("_extracted_at", F.current_timestamp())

    if config.source_file and is_file_source and source_path:
        df = df.withColumn("_source_file", F.lit(source_path))

    if config.source_connection and source_connection:
        df = df.withColumn("_source_connection", F.lit(source_connection))

    if config.source_table and source_table:
        df = df.withColumn("_source_table", F.lit(source_table))

    return df

anonymize(df, columns, method, salt=None)

Anonymize columns using Spark functions.

Source code in odibi/engine/spark_engine.py
def anonymize(self, df, columns: List[str], method: str, salt: Optional[str] = None):
    """Anonymize columns using Spark functions."""
    from pyspark.sql.functions import col, concat, lit, regexp_replace, sha2

    ctx = get_logging_context().with_context(engine="spark")
    ctx.debug(
        "Anonymizing columns",
        columns=columns,
        method=method,
        has_salt=salt is not None,
    )

    res = df
    for c in columns:
        if c not in df.columns:
            ctx.warning(f"Column '{c}' not found for anonymization, skipping", column=c)
            continue

        # Escape column names with backticks to handle special characters
        escaped_col = f"`{c}`"
        if method == "hash":
            if salt:
                res = res.withColumn(c, sha2(concat(col(escaped_col), lit(salt)), 256))
            else:
                res = res.withColumn(c, sha2(col(escaped_col), 256))

        elif method == "mask":
            res = res.withColumn(c, regexp_replace(col(escaped_col), ".(?=.{4})", "*"))

        elif method == "redact":
            res = res.withColumn(c, lit("[REDACTED]"))

    ctx.debug(f"Anonymization completed using {method}")
    return res

count_nulls(df, columns)

Count nulls in specified columns.

Source code in odibi/engine/spark_engine.py
def count_nulls(self, df, columns: List[str]) -> Dict[str, int]:
    """Count nulls in specified columns."""
    from pyspark.sql.functions import col, count, when

    missing = set(columns) - set(df.columns)
    if missing:
        raise ValueError(f"Columns not found in DataFrame: {', '.join(missing)}")

    # Escape column names with backticks to handle special characters (., ,, spaces)
    aggs = [count(when(col(f"`{c}`").isNull(), c)).alias(c) for c in columns]
    rows = df.select(*aggs).collect()
    return rows[0].asDict() if rows else {c: 0 for c in columns}

count_rows(df)

Count rows in DataFrame.

Source code in odibi/engine/spark_engine.py
def count_rows(self, df) -> int:
    """Count rows in DataFrame."""
    return df.count()

execute_operation(operation, params, df)

Execute built-in operation on Spark DataFrame.

Source code in odibi/engine/spark_engine.py
def execute_operation(self, operation: str, params: Dict[str, Any], df) -> Any:
    """Execute built-in operation on Spark DataFrame."""
    ctx = get_logging_context().with_context(engine="spark")
    params = params or {}

    ctx.debug(f"Executing operation: {operation}", params=list(params.keys()))

    if operation == "pivot":
        group_by = params.get("group_by", [])
        pivot_column = params.get("pivot_column")
        value_column = params.get("value_column")
        agg_func = params.get("agg_func", "first")

        if not pivot_column or not value_column:
            ctx.error("Pivot requires 'pivot_column' and 'value_column'")
            raise ValueError("Pivot requires 'pivot_column' and 'value_column'")

        if isinstance(group_by, str):
            group_by = [group_by]

        agg_expr = {value_column: agg_func}
        return df.groupBy(*group_by).pivot(pivot_column).agg(agg_expr)

    elif operation == "drop_duplicates":
        subset = params.get("subset")
        if subset:
            if isinstance(subset, str):
                subset = [subset]
            return df.dropDuplicates(subset=subset)
        return df.dropDuplicates()

    elif operation == "fillna":
        value = params.get("value")
        subset = params.get("subset")
        return df.fillna(value, subset=subset)

    elif operation == "drop":
        columns = params.get("columns")
        if not columns:
            return df
        if isinstance(columns, str):
            columns = [columns]
        return df.drop(*columns)

    elif operation == "rename":
        columns = params.get("columns")
        if not columns:
            return df

        res = df
        for old_name, new_name in columns.items():
            res = res.withColumnRenamed(old_name, new_name)
        return res

    elif operation == "sort":
        by = params.get("by")
        ascending = params.get("ascending", True)

        if not by:
            return df

        if isinstance(by, str):
            by = [by]

        if not ascending:
            from pyspark.sql.functions import desc

            sort_cols = [desc(c) for c in by]
            return df.orderBy(*sort_cols)

        return df.orderBy(*by)

    elif operation == "sample":
        fraction = params.get("frac", 0.1)
        seed = params.get("random_state")
        with_replacement = params.get("replace", False)
        return df.sample(withReplacement=with_replacement, fraction=fraction, seed=seed)

    else:
        # Fallback: check if operation is a registered transformer
        from odibi.context import EngineContext
        from odibi.registry import FunctionRegistry

        ctx.debug(
            f"Checking registry for operation: {operation}",
            registered_functions=list(FunctionRegistry._functions.keys())[:10],
            has_function=FunctionRegistry.has_function(operation),
        )

        if FunctionRegistry.has_function(operation):
            ctx.debug(f"Executing registered transformer as operation: {operation}")
            func = FunctionRegistry.get_function(operation)
            param_model = FunctionRegistry.get_param_model(operation)

            # Create EngineContext from current df
            from odibi.context import SparkContext

            engine_ctx = EngineContext(
                context=SparkContext(self.spark),
                df=df,
                engine=self,
                engine_type=self.engine_type,
            )

            # Validate and instantiate params
            if param_model:
                validated_params = param_model(**params)
                result_ctx = func(engine_ctx, validated_params)
            else:
                result_ctx = func(engine_ctx, **params)

            return result_ctx.df

        ctx.error(f"Unsupported operation for Spark engine: {operation}")
        raise ValueError(f"Unsupported operation for Spark engine: {operation}")

execute_sql(sql, context=None)

Execute SQL query in Spark.

Parameters:

Name Type Description Default
sql str

SQL query string

required
context Any

Execution context (optional, not used for Spark)

None

Returns:

Type Description
Any

Spark DataFrame with query results

Source code in odibi/engine/spark_engine.py
def execute_sql(self, sql: str, context: Any = None) -> Any:
    """Execute SQL query in Spark.

    Args:
        sql: SQL query string
        context: Execution context (optional, not used for Spark)

    Returns:
        Spark DataFrame with query results
    """
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()

    ctx.debug("Executing Spark SQL", query_preview=sql[:200] if len(sql) > 200 else sql)

    try:
        result = self.spark.sql(sql)
        elapsed = (time.time() - start_time) * 1000
        partition_count = self._safe_partition_count(result)

        ctx.log_spark_metrics(partition_count=partition_count)
        ctx.info(
            "Spark SQL executed",
            elapsed_ms=round(elapsed, 2),
            partitions=partition_count,
        )

        return result

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        error_type = type(e).__name__
        clean_message = _extract_spark_error_message(e)

        if "AnalysisException" in error_type:
            ctx.error(
                "Spark SQL Analysis Error",
                error_type=error_type,
                error_message=clean_message,
                query_preview=sql[:200] if len(sql) > 200 else sql,
                elapsed_ms=round(elapsed, 2),
            )
            raise TransformError(f"Spark SQL Analysis Error: {clean_message}")

        if "ParseException" in error_type:
            ctx.error(
                "Spark SQL Parse Error",
                error_type=error_type,
                error_message=clean_message,
                query_preview=sql[:200] if len(sql) > 200 else sql,
                elapsed_ms=round(elapsed, 2),
            )
            raise TransformError(f"Spark SQL Parse Error: {clean_message}")

        ctx.error(
            "Spark SQL execution failed",
            error_type=error_type,
            error_message=clean_message,
            elapsed_ms=round(elapsed, 2),
        )
        raise TransformError(f"Spark SQL Error: {clean_message}")

filter_coalesce(df, col1, col2, op, value)

Filter using COALESCE(col1, col2) op value.

Automatically casts string columns to timestamp for proper comparison. Tries multiple date formats including Oracle-style (DD-MON-YY).

Source code in odibi/engine/spark_engine.py
def filter_coalesce(self, df, col1: str, col2: str, op: str, value: Any) -> Any:
    """Filter using COALESCE(col1, col2) op value.

    Automatically casts string columns to timestamp for proper comparison.
    Tries multiple date formats including Oracle-style (DD-MON-YY).
    """
    from pyspark.sql import functions as F
    from pyspark.sql.types import StringType

    # Escape column names with backticks to handle special characters
    escaped_col1 = f"`{col1}`"
    escaped_col2 = f"`{col2}`"
    col1_type = df.schema[col1].dataType
    col2_type = df.schema[col2].dataType

    if isinstance(col1_type, StringType):
        c1 = self._parse_string_to_timestamp(F.col(escaped_col1))
    else:
        c1 = F.col(escaped_col1)

    if isinstance(col2_type, StringType):
        c2 = self._parse_string_to_timestamp(F.col(escaped_col2))
    else:
        c2 = F.col(escaped_col2)

    coalesced = F.coalesce(c1, c2)

    if op == ">":
        return df.filter(coalesced > value)
    elif op == ">=":
        return df.filter(coalesced >= value)
    elif op == "<":
        return df.filter(coalesced < value)
    elif op == "<=":
        return df.filter(coalesced <= value)
    elif op == "=":
        return df.filter(coalesced == value)
    else:
        return df.filter(f"COALESCE({col1}, {col2}) {op} '{value}'")

filter_greater_than(df, column, value)

Filter DataFrame where column > value.

Automatically casts string columns to timestamp for proper comparison. Tries multiple date formats including Oracle-style (DD-MON-YY).

Source code in odibi/engine/spark_engine.py
def filter_greater_than(self, df, column: str, value: Any) -> Any:
    """Filter DataFrame where column > value.

    Automatically casts string columns to timestamp for proper comparison.
    Tries multiple date formats including Oracle-style (DD-MON-YY).
    """
    from pyspark.sql import functions as F
    from pyspark.sql.types import StringType

    # Escape column names with backticks to handle special characters
    escaped_col = f"`{column}`"
    col_type = df.schema[column].dataType
    if isinstance(col_type, StringType):
        ts_col = self._parse_string_to_timestamp(F.col(escaped_col))
        return df.filter(ts_col > value)
    return df.filter(F.col(escaped_col) > value)

get_delta_history(connection, path, limit=None)

Get Delta table history.

Source code in odibi/engine/spark_engine.py
def get_delta_history(
    self, connection: Any, path: str, limit: Optional[int] = None
) -> List[Dict[str, Any]]:
    """Get Delta table history."""
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()

    ctx.debug("Fetching Delta history", path=path, limit=limit)

    try:
        from delta.tables import DeltaTable
    except ImportError:
        ctx.error("Delta Lake support requires 'delta-spark'")
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[spark]' "
            "with delta-spark. "
            "See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    try:
        delta_table = DeltaTable.forPath(self.spark, full_path)
        history_df = delta_table.history(limit) if limit else delta_table.history()
        history = [row.asDict() for row in history_df.collect()]

        elapsed = (time.time() - start_time) * 1000
        ctx.info(
            "Delta history retrieved",
            path=path,
            versions_returned=len(history),
            elapsed_ms=round(elapsed, 2),
        )

        return history

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        ctx.error(
            "Failed to get Delta history",
            path=path,
            error_type=type(e).__name__,
            error_message=str(e),
            elapsed_ms=round(elapsed, 2),
        )
        raise

get_sample(df, n=10)

Get sample rows as list of dictionaries.

Source code in odibi/engine/spark_engine.py
def get_sample(self, df, n: int = 10) -> List[Dict[str, Any]]:
    """Get sample rows as list of dictionaries."""
    return [row.asDict() for row in df.limit(n).collect()]

get_schema(df)

Get DataFrame schema with types.

Source code in odibi/engine/spark_engine.py
def get_schema(self, df) -> Dict[str, str]:
    """Get DataFrame schema with types."""
    return {f.name: f.dataType.simpleString() for f in df.schema}

get_shape(df)

Get DataFrame shape as (rows, columns).

Source code in odibi/engine/spark_engine.py
def get_shape(self, df) -> Tuple[int, int]:
    """Get DataFrame shape as (rows, columns)."""
    return (df.count(), len(df.columns))

get_source_files(df)

Get list of source files that generated this DataFrame.

Source code in odibi/engine/spark_engine.py
def get_source_files(self, df) -> List[str]:
    """Get list of source files that generated this DataFrame."""
    try:
        return df.inputFiles()
    except Exception:
        return []

get_table_schema(connection, table=None, path=None, format=None)

Get schema of an existing table/file.

Source code in odibi/engine/spark_engine.py
def get_table_schema(
    self,
    connection: Any,
    table: Optional[str] = None,
    path: Optional[str] = None,
    format: Optional[str] = None,
) -> Optional[Dict[str, str]]:
    """Get schema of an existing table/file."""
    ctx = get_logging_context().with_context(engine="spark")

    try:
        if table:
            if self.spark.catalog.tableExists(table):
                schema = self.get_schema(self.spark.table(table))
                ctx.debug(f"Retrieved schema for table: {table}", columns=len(schema))
                return schema
        elif path:
            full_path = connection.get_path(path)
            if format == "delta":
                from delta.tables import DeltaTable

                if DeltaTable.isDeltaTable(self.spark, full_path):
                    schema = self.get_schema(DeltaTable.forPath(self.spark, full_path).toDF())
                    ctx.debug(f"Retrieved Delta schema: {path}", columns=len(schema))
                    return schema
            elif format == "parquet":
                schema = self.get_schema(self.spark.read.parquet(full_path))
                ctx.debug(f"Retrieved Parquet schema: {path}", columns=len(schema))
                return schema
            elif format:
                schema = self.get_schema(self.spark.read.format(format).load(full_path))
                ctx.debug(f"Retrieved schema: {path}", format=format, columns=len(schema))
                return schema
    except Exception as e:
        ctx.warning(
            "Failed to get schema",
            table=table,
            path=path,
            error_message=str(e),
        )
    return None

harmonize_schema(df, target_schema, policy)

Harmonize DataFrame schema with target schema according to policy.

Source code in odibi/engine/spark_engine.py
def harmonize_schema(self, df, target_schema: Dict[str, str], policy: Any):
    """Harmonize DataFrame schema with target schema according to policy."""
    from pyspark.sql.functions import col, lit

    from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

    ctx = get_logging_context().with_context(engine="spark")

    target_cols = list(target_schema.keys())
    current_cols = df.columns

    missing = set(target_cols) - set(current_cols)
    new_cols = set(current_cols) - set(target_cols)

    ctx.debug(
        "Schema harmonization",
        target_columns=len(target_cols),
        current_columns=len(current_cols),
        missing_columns=list(missing) if missing else None,
        new_columns=list(new_cols) if new_cols else None,
        policy_mode=policy.mode.value if hasattr(policy.mode, "value") else str(policy.mode),
    )

    # Check Validations
    if missing and policy.on_missing_columns == OnMissingColumns.FAIL:
        ctx.error(
            f"Schema Policy Violation: Missing columns {missing}",
            missing_columns=list(missing),
        )
        raise ValueError(f"Schema Policy Violation: Missing columns {missing}")

    if new_cols and policy.on_new_columns == OnNewColumns.FAIL:
        ctx.error(
            f"Schema Policy Violation: New columns {new_cols}",
            new_columns=list(new_cols),
        )
        raise ValueError(f"Schema Policy Violation: New columns {new_cols}")

    # Apply Transformations
    if policy.mode == SchemaMode.EVOLVE and policy.on_new_columns == OnNewColumns.ADD_NULLABLE:
        res = df
        for c in missing:
            res = res.withColumn(c, lit(None))
        ctx.debug("Schema evolved: added missing columns as null")
        return res
    else:
        select_exprs = []
        for c in target_cols:
            if c in current_cols:
                # Escape column names with backticks to handle special characters
                select_exprs.append(col(f"`{c}`"))
            else:
                select_exprs.append(lit(None).alias(c))

        ctx.debug("Schema enforced: projected to target schema")
        return df.select(*select_exprs)

maintain_table(connection, format, table=None, path=None, config=None)

Run table maintenance operations (optimize, vacuum).

Source code in odibi/engine/spark_engine.py
def maintain_table(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    config: Optional[Any] = None,
) -> None:
    """Run table maintenance operations (optimize, vacuum)."""
    if format != "delta" or not config or not config.enabled:
        return

    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()

    if table:
        target = table
    elif path:
        full_path = connection.get_path(path)
        target = f"delta.`{full_path}`"
    else:
        return

    ctx.debug("Starting table maintenance", target=target)

    try:
        ctx.debug(f"Running OPTIMIZE on {target}")
        self.spark.sql(f"OPTIMIZE {target}")

        retention = config.vacuum_retention_hours
        if retention is not None and retention > 0:
            ctx.debug(f"Running VACUUM on {target}", retention_hours=retention)
            self.spark.sql(f"VACUUM {target} RETAIN {retention} HOURS")

        elapsed = (time.time() - start_time) * 1000
        ctx.info(
            "Table maintenance completed",
            target=target,
            vacuum_retention_hours=retention,
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        ctx.warning(
            f"Auto-optimize failed for {target}",
            error_type=type(e).__name__,
            error_message=str(e),
            elapsed_ms=round(elapsed, 2),
        )

profile_nulls(df)

Calculate null percentage for each column.

Source code in odibi/engine/spark_engine.py
def profile_nulls(self, df) -> Dict[str, float]:
    """Calculate null percentage for each column."""
    from pyspark.sql.functions import col, mean, when

    aggs = []
    for c in df.columns:
        # Escape column names with backticks to handle special characters (., ,, spaces)
        escaped_col = f"`{c}`"
        aggs.append(mean(when(col(escaped_col).isNull(), 1).otherwise(0)).alias(c))

    if not aggs:
        return {}

    try:
        rows = df.select(*aggs).collect()
        return rows[0].asDict() if rows else {}
    except Exception:
        return {}

read(connection, format, table=None, path=None, streaming=False, schema=None, options=None, as_of_version=None, as_of_timestamp=None)

Read data using Spark.

Parameters:

Name Type Description Default
connection Any

Connection object (with get_path method)

required
format str

Data format (csv, parquet, json, delta, sql_server)

required
table Optional[str]

Table name

None
path Optional[str]

File path

None
streaming bool

Whether to read as a stream (readStream)

False
schema Optional[str]

Schema string in DDL format (required for streaming file sources)

None
options Optional[Dict[str, Any]]

Format-specific options (including versionAsOf for Delta time travel)

None
as_of_version Optional[int]

Time travel version

None
as_of_timestamp Optional[str]

Time travel timestamp

None

Returns:

Type Description
Any

Spark DataFrame (or Streaming DataFrame)

Source code in odibi/engine/spark_engine.py
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
def read(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    streaming: bool = False,
    schema: Optional[str] = None,
    options: Optional[Dict[str, Any]] = None,
    as_of_version: Optional[int] = None,
    as_of_timestamp: Optional[str] = None,
) -> Any:
    """Read data using Spark.

    Args:
        connection: Connection object (with get_path method)
        format: Data format (csv, parquet, json, delta, sql_server)
        table: Table name
        path: File path
        streaming: Whether to read as a stream (readStream)
        schema: Schema string in DDL format (required for streaming file sources)
        options: Format-specific options (including versionAsOf for Delta time travel)
        as_of_version: Time travel version
        as_of_timestamp: Time travel timestamp

    Returns:
        Spark DataFrame (or Streaming DataFrame)
    """
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()
    options = options or {}

    source_identifier = table or path or "unknown"
    ctx.debug(
        "Starting Spark read",
        format=format,
        source=source_identifier,
        streaming=streaming,
        as_of_version=as_of_version,
        as_of_timestamp=as_of_timestamp,
    )

    # Handle Time Travel options
    if as_of_version is not None:
        options["versionAsOf"] = as_of_version
        ctx.debug(f"Time travel enabled: version {as_of_version}")
    if as_of_timestamp is not None:
        options["timestampAsOf"] = as_of_timestamp
        ctx.debug(f"Time travel enabled: timestamp {as_of_timestamp}")

    # SQL Server / Azure SQL Support
    if is_sql_format(format):
        if streaming:
            ctx.error("Streaming not supported for SQL Server / Azure SQL")
            raise ValueError("Streaming not supported for SQL Server / Azure SQL yet.")

        if not hasattr(connection, "get_spark_options"):
            conn_type = type(connection).__name__
            msg = f"Connection type '{conn_type}' does not support Spark SQL read"
            ctx.error(msg, connection_type=conn_type)
            raise ValueError(msg)

        jdbc_options = connection.get_spark_options()
        merged_options = {**jdbc_options, **options}

        # Extract filter for SQL pushdown
        sql_filter = merged_options.pop("filter", None)

        if "query" in merged_options:
            merged_options.pop("dbtable", None)
            # If filter provided with query, append to WHERE clause
            if sql_filter:
                existing_query = merged_options["query"]
                # Wrap existing query and add filter
                if "WHERE" in existing_query.upper():
                    merged_options["query"] = f"({existing_query}) AND ({sql_filter})"
                else:
                    subquery = f"SELECT * FROM ({existing_query}) AS _subq WHERE {sql_filter}"
                    merged_options["query"] = subquery
                ctx.debug(f"Applied SQL pushdown filter to query: {sql_filter}")
        elif table:
            # Build query with filter pushdown instead of using dbtable
            if sql_filter:
                merged_options.pop("dbtable", None)
                merged_options["query"] = f"SELECT * FROM {table} WHERE {sql_filter}"
                ctx.debug(f"Applied SQL pushdown filter: {sql_filter}")
            else:
                merged_options["dbtable"] = table
        elif "dbtable" not in merged_options:
            ctx.error("SQL format requires 'table' config or 'query' option")
            raise ValueError("SQL format requires 'table' config or 'query' option")

        ctx.debug("Executing JDBC read", has_query="query" in merged_options)

        try:
            df = self.spark.read.format("jdbc").options(**merged_options).load()
            elapsed = (time.time() - start_time) * 1000
            partition_count = self._safe_partition_count(df)

            ctx.log_file_io(path=source_identifier, format=format, mode="read")
            ctx.log_spark_metrics(partition_count=partition_count)
            ctx.info(
                "JDBC read completed",
                source=source_identifier,
                elapsed_ms=round(elapsed, 2),
                partitions=partition_count,
            )
            return df

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "JDBC read failed",
                source=source_identifier,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    # Simulation format: delegate to Pandas, convert to Spark
    if format == "simulation":
        if streaming:
            ctx.error("Streaming not supported for simulation format")
            raise ValueError("Streaming not supported for simulation format")

        ctx.debug("Reading simulation via Pandas engine, converting to Spark")
        from odibi.engine.pandas_engine import PandasEngine

        # Use Pandas engine for simulation
        pandas_engine = PandasEngine()
        pdf = pandas_engine._read_simulation(options, ctx)

        # Convert Pandas DataFrame to Spark DataFrame
        df = self.spark.createDataFrame(pdf)
        elapsed = (time.time() - start_time) * 1000

        # Preserve HWM metadata in Spark DataFrame metadata
        if hasattr(pdf, "attrs") and "_simulation_max_timestamp" in pdf.attrs:
            # Note: Spark doesn't have attrs, so we'll attach as a property
            df._simulation_max_timestamp = pdf.attrs["_simulation_max_timestamp"]

        # Preserve random walk state
        if hasattr(pdf, "attrs") and "_simulation_random_walk_state" in pdf.attrs:
            df._simulation_random_walk_state = pdf.attrs["_simulation_random_walk_state"]

        # Preserve scheduled event state
        if hasattr(pdf, "attrs") and "_simulation_scheduled_event_state" in pdf.attrs:
            df._simulation_scheduled_event_state = pdf.attrs[
                "_simulation_scheduled_event_state"
            ]

        ctx.info(
            "Simulation read completed (via Pandas)",
            elapsed_ms=round(elapsed, 2),
            row_count=len(pdf),
        )
        return df

    # Read based on format
    if table:
        # Managed/External Table (Catalog)
        ctx.debug(f"Reading from catalog table: {table}")

        if streaming:
            reader = self.spark.readStream.format(format)
        else:
            reader = self.spark.read.format(format)

        for key, value in options.items():
            reader = reader.option(key, value)

        try:
            df = reader.table(table)

            if "filter" in options:
                df = df.filter(options["filter"])
                ctx.debug(f"Applied filter: {options['filter']}")

            elapsed = (time.time() - start_time) * 1000

            if not streaming:
                partition_count = self._safe_partition_count(df)
                ctx.log_spark_metrics(partition_count=partition_count)
                ctx.log_file_io(path=table, format=format, mode="read")
                ctx.info(
                    f"Table read completed: {table}",
                    elapsed_ms=round(elapsed, 2),
                    partitions=partition_count,
                )
            else:
                ctx.info(f"Streaming read started: {table}", elapsed_ms=round(elapsed, 2))

            return df

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                f"Table read failed: {table}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    elif path:
        # File Path
        full_path = connection.get_path(path)
        ctx.debug(f"Reading from path: {full_path}")

        # Excel format: delegate to Pandas, convert to Spark
        if format == "excel":
            if streaming:
                ctx.error("Streaming not supported for Excel format")
                raise ValueError("Streaming not supported for Excel format")

            ctx.debug("Reading Excel via Pandas engine (best Excel support)")
            from odibi.engine.pandas_engine import PandasEngine

            # Get storage_options from connection for Azure/cloud authentication
            storage_options = None
            if hasattr(connection, "pandas_storage_options"):
                storage_options = connection.pandas_storage_options()

            # Use Pandas engine for Excel reading
            pandas_engine = PandasEngine()
            pdf = pandas_engine._read_excel_with_patterns(
                full_path,
                sheet_pattern=options.pop("sheet_pattern", None),
                sheet_pattern_case_sensitive=options.pop("sheet_pattern_case_sensitive", False),
                add_source_file=options.pop("add_source_file", False),
                is_glob="*" in str(full_path) or "?" in str(full_path),
                ctx=ctx,
                storage_options=storage_options,
                **options,
            )

            # Convert Pandas DataFrame to Spark DataFrame
            df = self.spark.createDataFrame(pdf)
            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                f"Excel read completed (via Pandas): {path}",
                elapsed_ms=round(elapsed, 2),
                row_count=len(pdf),
            )
            return df

        # API format: delegate to ApiFetcher, convert to Spark
        if format == "api":
            if streaming:
                ctx.error("Streaming not supported for API format")
                raise ValueError("Streaming not supported for API format")

            from odibi.connections.api_fetcher import create_api_fetcher
            from odibi.connections.http import HttpConnection

            ctx.debug("Reading API via ApiFetcher", endpoint=str(path))

            if not isinstance(connection, HttpConnection):
                raise ValueError(
                    f"Cannot read API data: connection type '{type(connection).__name__}' "
                    "is not an HttpConnection. Use an HTTP connection for API format."
                )

            # Extract API-specific options
            api_options = options.copy()
            params = api_options.pop("params", {})
            max_records = api_options.pop("max_records", None)
            method = api_options.pop("method", "GET")
            request_body = api_options.pop("request_body", None)

            # Create fetcher with connection's base_url and headers
            fetcher = create_api_fetcher(
                base_url=connection.base_url,
                headers=connection.headers,
                options=api_options,
            )

            # Fetch data as Pandas DataFrame
            pdf = fetcher.fetch_dataframe(
                endpoint=str(full_path),
                params=params,
                max_records=max_records,
                method=method,
                request_body=request_body,
            )

            # Convert Pandas DataFrame to Spark DataFrame
            df = self.spark.createDataFrame(pdf)
            elapsed = (time.time() - start_time) * 1000
            ctx.info(
                f"API read completed: {path}",
                elapsed_ms=round(elapsed, 2),
                row_count=len(pdf),
            )
            return df

        # Auto-detect encoding for CSV (Batch only)
        if not streaming and format == "csv" and options.get("auto_encoding"):
            options = options.copy()
            options.pop("auto_encoding")

            if "encoding" not in options:
                try:
                    from odibi.utils.encoding import detect_encoding

                    detected = detect_encoding(connection, path)
                    if detected:
                        options["encoding"] = detected
                        ctx.debug(f"Detected encoding: {detected}", path=path)
                except ImportError:
                    pass
                except Exception as e:
                    ctx.warning(
                        f"Encoding detection failed for {path}",
                        error_message=str(e),
                    )

        # Resolve cloudFiles.schemaLocation through read connection if still relative
        # (Node level resolves via write connection first if available)
        if "cloudFiles.schemaLocation" in options:
            schema_location = options["cloudFiles.schemaLocation"]
            if not schema_location.startswith(
                ("abfss://", "s3://", "gs://", "dbfs://", "hdfs://", "wasbs://")
            ):
                options = options.copy()
                options["cloudFiles.schemaLocation"] = connection.get_path(schema_location)
                ctx.debug(
                    "Resolved cloudFiles.schemaLocation through read connection (fallback)",
                    original=schema_location,
                    resolved=options["cloudFiles.schemaLocation"],
                )

        if streaming:
            reader = self.spark.readStream.format(format)
            if schema:
                reader = reader.schema(schema)
                ctx.debug(f"Applied schema for streaming read: {schema[:100]}...")
            else:
                # Determine if we should warn about missing schema
                # Formats that can infer schema: delta, parquet, avro (embedded schema)
                # cloudFiles with schemaLocation or self-describing formats (avro, parquet) are fine
                should_warn = True

                if format in ["delta", "parquet"]:
                    should_warn = False
                elif format == "cloudFiles":
                    cloud_format = options.get("cloudFiles.format", "")
                    has_schema_location = "cloudFiles.schemaLocation" in options
                    # avro and parquet have embedded schemas
                    if cloud_format in ["avro", "parquet"] or has_schema_location:
                        should_warn = False

                if should_warn:
                    ctx.warning(
                        f"Streaming read from '{format}' format without schema. "
                        "Schema inference is not supported for streaming sources. "
                        "Consider adding 'schema' to your read config."
                    )
        else:
            reader = self.spark.read.format(format)
            if schema:
                reader = reader.schema(schema)

        for key, value in options.items():
            if key == "header" and isinstance(value, bool):
                value = str(value).lower()
            reader = reader.option(key, value)

        try:
            df = reader.load(full_path)

            if "filter" in options:
                df = df.filter(options["filter"])
                ctx.debug(f"Applied filter: {options['filter']}")

            elapsed = (time.time() - start_time) * 1000

            if not streaming:
                partition_count = self._safe_partition_count(df)
                ctx.log_spark_metrics(partition_count=partition_count)
                ctx.log_file_io(path=path, format=format, mode="read")
                ctx.info(
                    f"File read completed: {path}",
                    elapsed_ms=round(elapsed, 2),
                    partitions=partition_count,
                    format=format,
                )
            else:
                ctx.info(f"Streaming read started: {path}", elapsed_ms=round(elapsed, 2))

            return df

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                f"File read failed: {path}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
                format=format,
            )
            raise
    else:
        ctx.error("Either path or table must be provided")
        raise ValueError("Either path or table must be provided")

restore_delta(connection, path, version)

Restore Delta table to a specific version.

Source code in odibi/engine/spark_engine.py
def restore_delta(self, connection: Any, path: str, version: int) -> None:
    """Restore Delta table to a specific version."""
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()

    ctx.debug("Restoring Delta table", path=path, version=version)

    try:
        from delta.tables import DeltaTable
    except ImportError:
        ctx.error("Delta Lake support requires 'delta-spark'")
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[spark]' "
            "with delta-spark. "
            "See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    try:
        delta_table = DeltaTable.forPath(self.spark, full_path)
        delta_table.restoreToVersion(version)

        elapsed = (time.time() - start_time) * 1000
        ctx.info(
            "Delta table restored",
            path=path,
            version=version,
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        ctx.error(
            "Delta restore failed",
            path=path,
            version=version,
            error_type=type(e).__name__,
            error_message=str(e),
            elapsed_ms=round(elapsed, 2),
        )
        raise

table_exists(connection, table=None, path=None)

Check if table or location exists.

Handles orphan catalog entries where the table is registered but the underlying Delta path no longer exists.

Source code in odibi/engine/spark_engine.py
def table_exists(
    self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
) -> bool:
    """Check if table or location exists.

    Handles orphan catalog entries where the table is registered but
    the underlying Delta path no longer exists.
    """
    ctx = get_logging_context().with_context(engine="spark")

    if table:
        try:
            if not self.spark.catalog.tableExists(table):
                ctx.debug(f"Table does not exist: {table}")
                return False
            # Table exists in catalog - verify it's actually readable
            # This catches orphan entries where path was deleted
            # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
            self.spark.table(table).limit(1).collect()
            ctx.debug(f"Table existence check: {table}", exists=True)
            return True
        except Exception as e:
            # Table exists in catalog but underlying data is gone (orphan entry)
            # This is expected during first-run detection - log at debug level
            ctx.debug(
                f"Table {table} exists in catalog but is not accessible (treating as first run)",
                error_message=str(e),
            )
            return False
    elif path:
        try:
            from delta.tables import DeltaTable

            full_path = connection.get_path(path)
            exists = DeltaTable.isDeltaTable(self.spark, full_path)
            ctx.debug(f"Delta table existence check: {path}", exists=exists)
            return exists
        except ImportError:
            try:
                full_path = connection.get_path(path)
                exists = (
                    self.spark.sparkContext._gateway.jvm.org.apache.hadoop.fs.FileSystem.get(
                        self.spark.sparkContext._jsc.hadoopConfiguration()
                    ).exists(
                        self.spark.sparkContext._gateway.jvm.org.apache.hadoop.fs.Path(
                            full_path
                        )
                    )
                )
                ctx.debug(f"Path existence check: {path}", exists=exists)
                return exists
            except Exception as e:
                ctx.warning(f"Path existence check failed: {path}", error_message=str(e))
                return False
        except Exception as e:
            ctx.warning(f"Table existence check failed: {path}", error_message=str(e))
            return False
    return False

vacuum_delta(connection, path, retention_hours=168)

VACUUM a Delta table to remove old files.

Source code in odibi/engine/spark_engine.py
def vacuum_delta(
    self,
    connection: Any,
    path: str,
    retention_hours: int = 168,
) -> None:
    """VACUUM a Delta table to remove old files."""
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()

    ctx.debug(
        "Starting Delta VACUUM",
        path=path,
        retention_hours=retention_hours,
    )

    try:
        from delta.tables import DeltaTable
    except ImportError:
        ctx.error("Delta Lake support requires 'delta-spark'")
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[spark]' "
            "with delta-spark. "
            "See README.md for installation instructions."
        )

    full_path = connection.get_path(path)

    try:
        delta_table = DeltaTable.forPath(self.spark, full_path)
        delta_table.vacuum(retention_hours / 24.0)

        elapsed = (time.time() - start_time) * 1000
        ctx.info(
            "Delta VACUUM completed",
            path=path,
            retention_hours=retention_hours,
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        ctx.error(
            "Delta VACUUM failed",
            path=path,
            error_type=type(e).__name__,
            error_message=str(e),
            elapsed_ms=round(elapsed, 2),
        )
        raise

validate_data(df, validation_config)

Validate DataFrame against rules.

Source code in odibi/engine/spark_engine.py
def validate_data(self, df, validation_config: Any) -> List[str]:
    """Validate DataFrame against rules."""
    from pyspark.sql.functions import col

    ctx = get_logging_context().with_context(engine="spark")
    failures = []

    if validation_config.not_empty:
        if df.isEmpty():
            failures.append("DataFrame is empty")

    if validation_config.no_nulls:
        null_counts = self.count_nulls(df, validation_config.no_nulls)
        for col_name, count in null_counts.items():
            if count > 0:
                failures.append(f"Column '{col_name}' has {count} null values")

    if validation_config.schema_validation:
        schema_failures = self.validate_schema(df, validation_config.schema_validation)
        failures.extend(schema_failures)

    if validation_config.ranges:
        for col_name, bounds in validation_config.ranges.items():
            if col_name in df.columns:
                min_val = bounds.get("min")
                max_val = bounds.get("max")

                # Escape column names with backticks to handle special characters
                escaped_col = f"`{col_name}`"
                if min_val is not None:
                    count = df.filter(col(escaped_col) < min_val).count()
                    if count > 0:
                        failures.append(f"Column '{col_name}' has values < {min_val}")

                if max_val is not None:
                    count = df.filter(col(escaped_col) > max_val).count()
                    if count > 0:
                        failures.append(f"Column '{col_name}' has values > {max_val}")
            else:
                failures.append(f"Column '{col_name}' not found for range validation")

    if validation_config.allowed_values:
        for col_name, allowed in validation_config.allowed_values.items():
            if col_name in df.columns:
                # Escape column names with backticks to handle special characters
                escaped_col = f"`{col_name}`"
                count = df.filter(~col(escaped_col).isin(allowed)).count()
                if count > 0:
                    failures.append(f"Column '{col_name}' has invalid values")
            else:
                failures.append(f"Column '{col_name}' not found for allowed values validation")

    ctx.log_validation_result(
        passed=len(failures) == 0,
        rule_name="data_validation",
        failures=failures if failures else None,
    )

    return failures

validate_schema(df, schema_rules)

Validate DataFrame schema.

Source code in odibi/engine/spark_engine.py
def validate_schema(self, df, schema_rules: Dict[str, Any]) -> List[str]:
    """Validate DataFrame schema."""
    failures = []

    if "required_columns" in schema_rules:
        required = schema_rules["required_columns"]
        missing = set(required) - set(df.columns)
        if missing:
            failures.append(f"Missing required columns: {', '.join(missing)}")

    if "types" in schema_rules:
        type_map = {
            "int": ["integer", "long", "short", "byte", "bigint"],
            "float": ["double", "float"],
            "str": ["string"],
            "bool": ["boolean"],
        }

        for col_name, expected_type in schema_rules["types"].items():
            if col_name not in df.columns:
                failures.append(f"Column '{col_name}' not found for type validation")
                continue

            actual_type = dict(df.dtypes)[col_name]
            expected_dtypes = type_map.get(expected_type, [expected_type])

            if actual_type not in expected_dtypes:
                failures.append(
                    f"Column '{col_name}' has type '{actual_type}', expected '{expected_type}'"
                )

    return failures

write(df, connection, format, table=None, path=None, register_table=None, mode='overwrite', options=None, streaming_config=None)

Write data using Spark.

Parameters:

Name Type Description Default
df Any

Spark DataFrame to write

required
connection Any

Connection object

required
format str

Output format (csv, parquet, json, delta)

required
table Optional[str]

Table name

None
path Optional[str]

File path

None
register_table Optional[str]

Name to register as external table (if path is used)

None
mode str

Write mode (overwrite, append, error, ignore, upsert, append_once)

'overwrite'
options Optional[Dict[str, Any]]

Format-specific options (including partition_by for partitioning)

None
streaming_config Optional[Any]

StreamingWriteConfig for streaming DataFrames

None

Returns:

Type Description
Optional[Dict[str, Any]]

Optional dictionary containing Delta commit metadata (if format=delta),

Optional[Dict[str, Any]]

or streaming query info (if streaming)

Source code in odibi/engine/spark_engine.py
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
def write(
    self,
    df: Any,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    register_table: Optional[str] = None,
    mode: str = "overwrite",
    options: Optional[Dict[str, Any]] = None,
    streaming_config: Optional[Any] = None,
) -> Optional[Dict[str, Any]]:
    """Write data using Spark.

    Args:
        df: Spark DataFrame to write
        connection: Connection object
        format: Output format (csv, parquet, json, delta)
        table: Table name
        path: File path
        register_table: Name to register as external table (if path is used)
        mode: Write mode (overwrite, append, error, ignore, upsert, append_once)
        options: Format-specific options (including partition_by for partitioning)
        streaming_config: StreamingWriteConfig for streaming DataFrames

    Returns:
        Optional dictionary containing Delta commit metadata (if format=delta),
        or streaming query info (if streaming)
    """
    ctx = get_logging_context().with_context(engine="spark")
    start_time = time.time()
    options = options or {}

    if getattr(df, "isStreaming", False) is True:
        return self._write_streaming(
            df=df,
            connection=connection,
            format=format,
            table=table,
            path=path,
            register_table=register_table,
            options=options,
            streaming_config=streaming_config,
        )

    target_identifier = table or path or "unknown"
    try:
        partition_count = df.rdd.getNumPartitions()
    except Exception:
        partition_count = 1  # Fallback for mocks or unsupported DataFrames

    # Auto-coalesce DataFrames for Delta writes to reduce file overhead
    # Use coalesce_partitions option to explicitly set target partitions
    # NOTE: We avoid df.count() here as it would trigger double-evaluation of lazy DataFrames
    coalesce_partitions = options.pop("coalesce_partitions", None)
    if (
        coalesce_partitions
        and isinstance(partition_count, int)
        and partition_count > coalesce_partitions
    ):
        df = df.coalesce(coalesce_partitions)
        ctx.debug(
            f"Coalesced DataFrame to {coalesce_partitions} partition(s)",
            original_partitions=partition_count,
        )
        partition_count = coalesce_partitions

    ctx.debug(
        "Starting Spark write",
        format=format,
        target=target_identifier,
        mode=mode,
        partitions=partition_count,
    )

    # SQL Server / Azure SQL Support
    if is_sql_format(format):
        if not hasattr(connection, "get_spark_options"):
            conn_type = type(connection).__name__
            msg = f"Connection type '{conn_type}' does not support Spark SQL write"
            ctx.error(msg, connection_type=conn_type)
            raise ValueError(msg)

        jdbc_options = connection.get_spark_options()
        merged_options = {**jdbc_options, **options}

        if table:
            merged_options["dbtable"] = table
        elif "dbtable" not in merged_options:
            ctx.error("SQL format requires 'table' config or 'dbtable' option")
            raise ValueError("SQL format requires 'table' config or 'dbtable' option")

        # Handle MERGE mode (SQL Server only)
        if mode == "merge":
            if getattr(connection, "sql_dialect", "mssql") != "mssql":
                raise NotImplementedError(
                    f"SQL MERGE mode is currently only supported for SQL Server connections. "
                    f"Connection dialect '{getattr(connection, 'sql_dialect', 'unknown')}' "
                    f"does not support MERGE. Use mode='append' or mode='overwrite' instead."
                )

            merge_keys = options.get("merge_keys")
            merge_options = options.get("merge_options")

            if not merge_keys:
                ctx.error("MERGE mode requires 'merge_keys' in options")
                raise ValueError(
                    "MERGE mode requires 'merge_keys' in options. "
                    "Specify the key columns for the MERGE ON clause."
                )

            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            writer = SqlServerMergeWriter(connection)

            # Check if bulk_copy is enabled for merge staging
            staging_connection = None
            if merge_options and getattr(merge_options, "bulk_copy", False):
                staging_conn_name = getattr(merge_options, "staging_connection", None)
                if not staging_conn_name:
                    raise ValueError(
                        "bulk_copy=True requires 'staging_connection' to be set. "
                        "Specify an ADLS/Blob connection for staging files."
                    )
                # Get staging connection from options (passed from node executor)
                staging_connection = options.get("_staging_connection")
                if not staging_connection:
                    raise ValueError(
                        f"Staging connection '{staging_conn_name}' not found. "
                        "Ensure the connection is defined in your project config."
                    )

            ctx.debug(
                "Executing SQL Server MERGE",
                target=table,
                merge_keys=merge_keys,
                bulk_copy=staging_connection is not None,
            )

            try:
                result = writer.merge(
                    df=df,
                    spark_engine=self,
                    target_table=table,
                    merge_keys=merge_keys,
                    options=merge_options,
                    jdbc_options=jdbc_options,
                    staging_connection=staging_connection,
                )
                elapsed = (time.time() - start_time) * 1000
                ctx.log_file_io(path=target_identifier, format=format, mode="write")
                ctx.info(
                    "SQL Server MERGE completed",
                    target=target_identifier,
                    mode=mode,
                    inserted=result.inserted,
                    updated=result.updated,
                    deleted=result.deleted,
                    elapsed_ms=round(elapsed, 2),
                )
                return {
                    "mode": "merge",
                    "inserted": result.inserted,
                    "updated": result.updated,
                    "deleted": result.deleted,
                    "total_affected": result.total_affected,
                }

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "SQL Server MERGE failed",
                    target=target_identifier,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        # Handle enhanced overwrite with strategies
        if mode == "overwrite" and options.get("overwrite_options"):
            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            overwrite_options = options.get("overwrite_options")
            writer = SqlServerMergeWriter(connection)

            # Check if bulk_copy is enabled
            use_bulk_copy = getattr(overwrite_options, "bulk_copy", False)

            if use_bulk_copy:
                # Bulk copy mode - use staging files + BULK INSERT
                staging_conn_name = getattr(overwrite_options, "staging_connection", None)
                external_data_source = getattr(overwrite_options, "external_data_source", None)

                if not staging_conn_name:
                    raise ValueError(
                        "bulk_copy=True requires 'staging_connection' to be set. "
                        "Specify an ADLS/Blob connection for staging files."
                    )
                if not external_data_source:
                    raise ValueError(
                        "bulk_copy=True requires 'external_data_source' to be set. "
                        "This is the SQL Server external data source name pointing "
                        "to your staging storage."
                    )

                # Get staging connection from options (passed from node executor)
                staging_connection = options.get("_staging_connection")
                if not staging_connection:
                    raise ValueError(
                        f"Staging connection '{staging_conn_name}' not found. "
                        "Ensure the connection is defined in your project config."
                    )

                ctx.debug(
                    "Executing SQL Server bulk copy overwrite",
                    target=table,
                    staging_connection=staging_conn_name,
                    external_data_source=external_data_source,
                )

                # Get bulk copy context for staging path organization
                bulk_copy_context = options.get("_bulk_copy_context", {})

                try:
                    result = writer.bulk_copy_spark(
                        df=df,
                        target_table=table,
                        staging_connection=staging_connection,
                        external_data_source=external_data_source,
                        options=overwrite_options,
                        bulk_copy_context=bulk_copy_context,
                    )
                    elapsed = (time.time() - start_time) * 1000
                    ctx.log_file_io(path=target_identifier, format=format, mode="write")
                    ctx.info(
                        "SQL Server bulk copy completed",
                        target=target_identifier,
                        strategy="bulk_copy",
                        rows_written=result.rows_written,
                        elapsed_ms=round(elapsed, 2),
                    )
                    return {
                        "mode": "overwrite",
                        "strategy": "bulk_copy",
                        "rows_written": result.rows_written,
                    }

                except Exception as e:
                    elapsed = (time.time() - start_time) * 1000
                    ctx.error(
                        "SQL Server bulk copy failed",
                        target=target_identifier,
                        error_type=type(e).__name__,
                        error_message=str(e),
                        elapsed_ms=round(elapsed, 2),
                    )
                    raise

            # Standard JDBC-based overwrite
            ctx.debug(
                "Executing SQL Server enhanced overwrite",
                target=table,
                strategy=(
                    overwrite_options.strategy.value
                    if hasattr(overwrite_options, "strategy")
                    else "truncate_insert"
                ),
            )

            try:
                result = writer.overwrite_spark(
                    df=df,
                    target_table=table,
                    options=overwrite_options,
                    jdbc_options=jdbc_options,
                )
                elapsed = (time.time() - start_time) * 1000
                ctx.log_file_io(path=target_identifier, format=format, mode="write")
                ctx.info(
                    "SQL Server enhanced overwrite completed",
                    target=target_identifier,
                    strategy=result.strategy,
                    rows_written=result.rows_written,
                    elapsed_ms=round(elapsed, 2),
                )
                return {
                    "mode": "overwrite",
                    "strategy": result.strategy,
                    "rows_written": result.rows_written,
                }

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "SQL Server enhanced overwrite failed",
                    target=target_identifier,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

        if mode not in ["overwrite", "append", "ignore", "error"]:
            if mode == "fail":
                mode = "error"
            else:
                ctx.error(f"Write mode '{mode}' not supported for Spark SQL write")
                raise ValueError(f"Write mode '{mode}' not supported for Spark SQL write")

        ctx.debug("Executing JDBC write", target=table or merged_options.get("dbtable"))

        try:
            df.write.format("jdbc").options(**merged_options).mode(mode).save()
            elapsed = (time.time() - start_time) * 1000
            ctx.log_file_io(path=target_identifier, format=format, mode="write")
            ctx.info(
                "JDBC write completed",
                target=target_identifier,
                mode=mode,
                elapsed_ms=round(elapsed, 2),
            )
            return None

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                "JDBC write failed",
                target=target_identifier,
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    # Unity Catalog: qualify bare table/path names via the connection so
    # the configured catalog.schema is honoured.  Without this, Spark
    # resolves bare names against the session default (workspace.default),
    # silently ignoring the YAML schema parameter.
    from odibi.connections.unity_catalog import UnityCatalogConnection

    if isinstance(connection, UnityCatalogConnection):
        if not table and path:
            resolved = connection.get_path(path)
            if not resolved.startswith("/"):
                table = resolved
                path = None
        elif table and "." not in table:
            table = connection.get_path(table)

    # Handle Upsert/AppendOnce (Delta Only)
    if mode in ["upsert", "append_once"]:
        if format != "delta":
            ctx.error(f"Mode '{mode}' only supported for Delta format")
            raise NotImplementedError(
                f"Mode '{mode}' only supported for Delta format in Spark engine."
            )

        keys = options.get("keys")
        if not keys:
            ctx.error(f"Mode '{mode}' requires 'keys' list in options")
            raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

        if isinstance(keys, str):
            keys = [keys]

        exists = self.table_exists(connection, table, path)
        ctx.debug("Table existence check for merge", target=target_identifier, exists=exists)

        if not exists:
            mode = "overwrite"
            ctx.debug("Target does not exist, falling back to overwrite mode")
        else:
            from delta.tables import DeltaTable

            target_dt = None
            target_name = ""
            is_table_target = False

            if table:
                target_dt = DeltaTable.forName(self.spark, table)
                target_name = table
                is_table_target = True
            elif path:
                full_path = connection.get_path(path)
                target_dt = DeltaTable.forPath(self.spark, full_path)
                target_name = full_path
                is_table_target = False

            condition = " AND ".join([f"target.`{k}` = source.`{k}`" for k in keys])
            ctx.debug("Executing Delta merge", merge_mode=mode, keys=keys, condition=condition)

            merge_builder = target_dt.alias("target").merge(df.alias("source"), condition)

            try:
                if mode == "upsert":
                    merge_builder.whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()
                elif mode == "append_once":
                    merge_builder.whenNotMatchedInsertAll().execute()

                elapsed = (time.time() - start_time) * 1000
                ctx.info(
                    "Delta merge completed",
                    target=target_name,
                    mode=mode,
                    elapsed_ms=round(elapsed, 2),
                )

                self._optimize_delta_write(target_name, options, is_table=is_table_target)
                commit_info = self._get_last_delta_commit_info(
                    target_name, is_table=is_table_target
                )

                if commit_info:
                    ctx.debug(
                        "Delta commit info",
                        version=commit_info.get("version"),
                        operation=commit_info.get("operation"),
                    )

                return commit_info

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "Delta merge failed",
                    target=target_name,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

    # Get output location
    if table:
        # Managed/External Table (Catalog)
        ctx.debug(f"Writing to catalog table: {table}")
        writer = df.write.format(format).mode(mode)

        partition_by = options.get("partition_by")
        if partition_by:
            if isinstance(partition_by, str):
                partition_by = [partition_by]
            writer = writer.partitionBy(*partition_by)
            ctx.debug(f"Partitioning by: {partition_by}")

        for key, value in options.items():
            writer = writer.option(key, value)

        try:
            writer.saveAsTable(table)
            elapsed = (time.time() - start_time) * 1000

            ctx.log_file_io(
                path=table,
                format=format,
                mode=mode,
                partitions=partition_by,
            )
            ctx.info(
                f"Table write completed: {table}",
                mode=mode,
                elapsed_ms=round(elapsed, 2),
            )

            if format == "delta":
                self._optimize_delta_write(table, options, is_table=True)
                return self._get_last_delta_commit_info(table, is_table=True)
            return None

        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            ctx.error(
                f"Table write failed: {table}",
                error_type=type(e).__name__,
                error_message=str(e),
                elapsed_ms=round(elapsed, 2),
            )
            raise

    elif path:
        full_path = connection.get_path(path)
    else:
        ctx.error("Either path or table must be provided")
        raise ValueError("Either path or table must be provided")

    # Extract partition_by option
    partition_by = options.pop("partition_by", None) or options.pop("partitionBy", None)

    # Extract cluster_by option (Liquid Clustering)
    cluster_by = options.pop("cluster_by", None)

    # Warn about partitioning anti-patterns
    if partition_by and cluster_by:
        import warnings

        ctx.warning(
            "Conflict: Both 'partition_by' and 'cluster_by' are set",
            partition_by=partition_by,
            cluster_by=cluster_by,
        )
        warnings.warn(
            "⚠️  Conflict: Both 'partition_by' and 'cluster_by' (Liquid Clustering) are set. "
            "Liquid Clustering supersedes partitioning. 'partition_by' will be ignored "
            "if the table is being created now.",
            UserWarning,
        )

    elif partition_by:
        import warnings

        ctx.warning(
            "Partitioning warning: ensure low-cardinality columns",
            partition_by=partition_by,
        )
        warnings.warn(
            "⚠️  Partitioning can cause performance issues if misused. "
            "Only partition on low-cardinality columns (< 1000 unique values) "
            "and ensure each partition has > 1000 rows.",
            UserWarning,
        )

    # Handle Upsert/Append-Once for Delta Lake (Path-based only for now)
    if format == "delta" and mode in ["upsert", "append_once"]:
        try:
            from delta.tables import DeltaTable
        except ImportError:
            ctx.error("Delta Lake support requires 'delta-spark'")
            raise ImportError("Delta Lake support requires 'delta-spark'")

        if "keys" not in options:
            ctx.error(f"Mode '{mode}' requires 'keys' list in options")
            raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

        if DeltaTable.isDeltaTable(self.spark, full_path):
            ctx.debug(f"Performing Delta merge at path: {full_path}")
            delta_table = DeltaTable.forPath(self.spark, full_path)
            keys = options["keys"]
            if isinstance(keys, str):
                keys = [keys]

            condition = " AND ".join([f"target.{k} = source.{k}" for k in keys])
            merger = delta_table.alias("target").merge(df.alias("source"), condition)

            try:
                if mode == "upsert":
                    merger.whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()
                else:
                    merger.whenNotMatchedInsertAll().execute()

                elapsed = (time.time() - start_time) * 1000
                ctx.info(
                    "Delta merge completed at path",
                    path=path,
                    mode=mode,
                    elapsed_ms=round(elapsed, 2),
                )

                if register_table:
                    try:
                        table_in_catalog = self.spark.catalog.tableExists(register_table)
                        needs_registration = not table_in_catalog

                        # Handle orphan catalog entries (only for path-not-found errors)
                        if table_in_catalog:
                            try:
                                # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
                                self.spark.table(register_table).limit(1).collect()
                                ctx.debug(
                                    f"Table '{register_table}' already registered and valid"
                                )
                            except Exception as verify_err:
                                error_str = str(verify_err)
                                is_orphan = (
                                    "DELTA_PATH_DOES_NOT_EXIST" in error_str
                                    or "Path does not exist" in error_str
                                    or "FileNotFoundException" in error_str
                                )
                                if is_orphan:
                                    ctx.warning(
                                        f"Table '{register_table}' is orphan, re-registering"
                                    )
                                    try:
                                        self.spark.sql(f"DROP TABLE IF EXISTS {register_table}")
                                    except Exception:
                                        pass
                                    needs_registration = True
                                else:
                                    ctx.debug(
                                        f"Table '{register_table}' verify failed, "
                                        "skipping registration"
                                    )

                        if needs_registration:
                            create_sql = (
                                f"CREATE TABLE IF NOT EXISTS {register_table} "
                                f"USING DELTA LOCATION '{full_path}'"
                            )
                            self.spark.sql(create_sql)
                            ctx.info(f"Registered table: {register_table}", path=full_path)
                    except Exception as e:
                        ctx.error(
                            f"Failed to register external table '{register_table}'",
                            error_message=str(e),
                        )

                self._optimize_delta_write(full_path, options, is_table=False)
                return self._get_last_delta_commit_info(full_path, is_table=False)

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "Delta merge failed at path",
                    path=path,
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise
        else:
            mode = "overwrite"
            ctx.debug("Target does not exist, falling back to overwrite mode")

    # Write based on format (Path-based)
    ctx.debug(f"Writing to path: {full_path}")

    # Handle Liquid Clustering (New Table Creation via SQL)
    if format == "delta" and cluster_by:
        should_create = False
        target_name = None

        if table:
            target_name = table
            if mode == "overwrite":
                should_create = True
            elif mode == "append":
                if not self.spark.catalog.tableExists(table):
                    should_create = True
        elif path:
            full_path = connection.get_path(path)
            target_name = f"delta.`{full_path}`"
            if mode == "overwrite":
                should_create = True
            elif mode == "append":
                try:
                    from delta.tables import DeltaTable

                    if not DeltaTable.isDeltaTable(self.spark, full_path):
                        should_create = True
                except ImportError:
                    pass

        if should_create:
            if isinstance(cluster_by, str):
                cluster_by = [cluster_by]

            cols = ", ".join(cluster_by)
            temp_view = f"odibi_temp_writer_{abs(hash(str(target_name)))}"
            df.createOrReplaceTempView(temp_view)

            create_cmd = (
                "CREATE OR REPLACE TABLE"
                if mode == "overwrite"
                else "CREATE TABLE IF NOT EXISTS"
            )

            sql = (
                f"{create_cmd} {target_name} USING DELTA CLUSTER BY ({cols}) "
                f"AS SELECT * FROM {temp_view}"
            )

            ctx.debug("Creating clustered Delta table", sql=sql, cluster_by=cluster_by)

            try:
                self.spark.sql(sql)
                self.spark.catalog.dropTempView(temp_view)

                elapsed = (time.time() - start_time) * 1000
                ctx.info(
                    "Clustered Delta table created",
                    target=target_name,
                    cluster_by=cluster_by,
                    elapsed_ms=round(elapsed, 2),
                )

                if register_table and path:
                    try:
                        reg_sql = (
                            f"CREATE TABLE IF NOT EXISTS {register_table} "
                            f"USING DELTA LOCATION '{full_path}'"
                        )
                        self.spark.sql(reg_sql)
                        ctx.info(f"Registered table: {register_table}")
                    except Exception:
                        pass

                if format == "delta":
                    self._optimize_delta_write(
                        target_name if table else full_path, options, is_table=bool(table)
                    )
                    return self._get_last_delta_commit_info(
                        target_name if table else full_path, is_table=bool(table)
                    )
                return None

            except Exception as e:
                elapsed = (time.time() - start_time) * 1000
                ctx.error(
                    "Failed to create clustered Delta table",
                    error_type=type(e).__name__,
                    error_message=str(e),
                    elapsed_ms=round(elapsed, 2),
                )
                raise

    # Extract table_properties from options
    table_properties = options.pop("table_properties", None)

    # For column mapping and other properties that must be set BEFORE write
    original_configs = {}
    if table_properties and format == "delta":
        for prop_name, prop_value in table_properties.items():
            spark_conf_key = (
                f"spark.databricks.delta.properties.defaults.{prop_name.replace('delta.', '')}"
            )
            try:
                original_configs[spark_conf_key] = self.spark.conf.get(spark_conf_key, None)
            except Exception:
                original_configs[spark_conf_key] = None
            self.spark.conf.set(spark_conf_key, prop_value)
        ctx.debug(
            "Applied table properties as session defaults",
            properties=list(table_properties.keys()),
        )

    writer = df.write.format(format).mode(mode)

    if partition_by:
        if isinstance(partition_by, str):
            partition_by = [partition_by]
        writer = writer.partitionBy(*partition_by)
        ctx.debug(f"Partitioning by: {partition_by}")

    for key, value in options.items():
        writer = writer.option(key, value)

    try:
        writer.save(full_path)
        elapsed = (time.time() - start_time) * 1000

        ctx.log_file_io(
            path=path,
            format=format,
            mode=mode,
            partitions=partition_by,
        )
        ctx.info(
            f"File write completed: {path}",
            format=format,
            mode=mode,
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        elapsed = (time.time() - start_time) * 1000
        ctx.error(
            f"File write failed: {path}",
            error_type=type(e).__name__,
            error_message=str(e),
            elapsed_ms=round(elapsed, 2),
        )
        raise
    finally:
        for conf_key, original_value in original_configs.items():
            if original_value is None:
                self.spark.conf.unset(conf_key)
            else:
                self.spark.conf.set(conf_key, original_value)

    if format == "delta":
        self._optimize_delta_write(full_path, options, is_table=False)

    if register_table and format == "delta":
        try:
            table_in_catalog = self.spark.catalog.tableExists(register_table)
            needs_registration = not table_in_catalog

            # Handle orphan catalog entries: table exists but points to deleted path
            # Only treat as orphan if it's specifically a DELTA_PATH_DOES_NOT_EXIST error
            if table_in_catalog:
                try:
                    # Use limit(1) not limit(0) - limit(0) can succeed from metadata alone
                    self.spark.table(register_table).limit(1).collect()
                    ctx.debug(
                        f"Table '{register_table}' already registered and valid, "
                        "skipping registration"
                    )
                except Exception as verify_err:
                    error_str = str(verify_err)
                    is_orphan = (
                        "DELTA_PATH_DOES_NOT_EXIST" in error_str
                        or "Path does not exist" in error_str
                        or "FileNotFoundException" in error_str
                    )

                    if is_orphan:
                        # Orphan entry - table in catalog but path was deleted
                        ctx.warning(
                            f"Table '{register_table}' is orphan (path deleted), "
                            "dropping and re-registering",
                            error_message=error_str[:200],
                        )
                        try:
                            self.spark.sql(f"DROP TABLE IF EXISTS {register_table}")
                        except Exception:
                            pass  # Best effort cleanup
                        needs_registration = True
                    else:
                        # Other error (auth, network, etc.) - don't drop, just log
                        ctx.debug(
                            f"Table '{register_table}' exists but verify failed "
                            "(not orphan), skipping registration",
                            error_message=error_str[:200],
                        )

            if needs_registration:
                ctx.debug(f"Registering table '{register_table}' at '{full_path}'")
                reg_sql = (
                    f"CREATE TABLE IF NOT EXISTS {register_table} "
                    f"USING DELTA LOCATION '{full_path}'"
                )
                self.spark.sql(reg_sql)
                ctx.info(f"Registered table: {register_table}", path=full_path)
        except Exception as e:
            ctx.error(
                f"Failed to register table '{register_table}'",
                error_message=str(e),
            )
            raise RuntimeError(
                f"Failed to register external table '{register_table}': {e}"
            ) from e

    if format == "delta":
        return self._get_last_delta_commit_info(full_path, is_table=False)

    return None

odibi.engine.polars_engine

Polars engine implementation.

PolarsEngine

Bases: Engine

Polars-based execution engine (High Performance).

Source code in odibi/engine/polars_engine.py
  25
  26
  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
  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
 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
 433
 434
 435
 436
 437
 438
 439
 440
 441
 442
 443
 444
 445
 446
 447
 448
 449
 450
 451
 452
 453
 454
 455
 456
 457
 458
 459
 460
 461
 462
 463
 464
 465
 466
 467
 468
 469
 470
 471
 472
 473
 474
 475
 476
 477
 478
 479
 480
 481
 482
 483
 484
 485
 486
 487
 488
 489
 490
 491
 492
 493
 494
 495
 496
 497
 498
 499
 500
 501
 502
 503
 504
 505
 506
 507
 508
 509
 510
 511
 512
 513
 514
 515
 516
 517
 518
 519
 520
 521
 522
 523
 524
 525
 526
 527
 528
 529
 530
 531
 532
 533
 534
 535
 536
 537
 538
 539
 540
 541
 542
 543
 544
 545
 546
 547
 548
 549
 550
 551
 552
 553
 554
 555
 556
 557
 558
 559
 560
 561
 562
 563
 564
 565
 566
 567
 568
 569
 570
 571
 572
 573
 574
 575
 576
 577
 578
 579
 580
 581
 582
 583
 584
 585
 586
 587
 588
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
class PolarsEngine(Engine):
    """Polars-based execution engine (High Performance)."""

    name = "polars"
    engine_type = EngineType.POLARS

    def __init__(
        self,
        connections: Optional[Dict[str, Any]] = None,
        config: Optional[Dict[str, Any]] = None,
    ):
        """Initialize Polars engine.

        Args:
            connections: Dictionary of connection objects
            config: Engine configuration (optional)
        """
        if pl is None:
            raise ImportError("Polars not installed. Run 'pip install polars'.")

        self.connections = connections or {}
        self.config = config or {}

    def materialize(self, df: Any) -> Any:
        """Materialize lazy dataset into memory (DataFrame).

        Args:
            df: LazyFrame or DataFrame

        Returns:
            Materialized DataFrame (pl.DataFrame)
        """
        if isinstance(df, pl.LazyFrame):
            return df.collect()
        return df

    def read(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        streaming: bool = False,
        schema: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Any:
        """Read data using Polars (Lazy by default).

        Args:
            connection: Connection object providing base path and storage options.
            format: Data format (csv, parquet, delta, json, sql, sql_server, azure_sql).
            table: Table name (mutually exclusive with path).
            path: File path (mutually exclusive with table).
            streaming: Whether to enable streaming mode (uses scan methods when possible).
            schema: Optional schema specification for SQL queries.
            options: Additional read options to pass to Polars readers.
            **kwargs: Additional keyword arguments.

        Returns:
            pl.LazyFrame or pl.DataFrame
        """
        options = options or {}

        # SQL Server / Azure SQL Support
        if is_sql_format(format):
            if not hasattr(connection, "read_table") and not hasattr(connection, "read_sql_query"):
                raise ValueError(
                    f"Cannot read SQL table: connection type '{type(connection).__name__}' "
                    "does not support SQL operations. Use a SQL-compatible connection."
                )
            table_name = table or path
            if not table_name:
                raise ValueError("SQL read requires 'table' or 'path' to specify the table name.")
            default_schema = getattr(connection, "default_schema", "dbo")
            if "." in table_name:
                schema_name, tbl = table_name.split(".", 1)
            else:
                schema_name, tbl = default_schema, table_name

            sql_filter = options.get("filter")

            if sql_filter and hasattr(connection, "read_sql_query"):
                full_query = connection.build_select_query(tbl, schema_name, where=sql_filter)
                pdf = connection.read_sql_query(full_query)
            else:
                pdf = connection.read_table(table_name=tbl, schema=schema_name)

            return pl.from_pandas(pdf).lazy()

        # Simulation format: delegate to Pandas, convert to Polars
        if format == "simulation":
            from odibi.engine.pandas_engine import PandasEngine
            from odibi.utils.logging_context import get_logging_context

            ctx = get_logging_context()
            ctx.debug("Reading simulation via Pandas engine, converting to Polars")

            # Use Pandas engine for simulation
            pandas_engine = PandasEngine()
            pdf = pandas_engine._read_simulation(options, ctx)

            # Convert to Polars LazyFrame
            lf = pl.from_pandas(pdf).lazy()

            # Preserve HWM metadata
            if hasattr(pdf, "attrs") and "_simulation_max_timestamp" in pdf.attrs:
                # Polars LazyFrame doesn't have attrs, so we attach as property
                lf._simulation_max_timestamp = pdf.attrs["_simulation_max_timestamp"]

            # Preserve random walk state
            if hasattr(pdf, "attrs") and "_simulation_random_walk_state" in pdf.attrs:
                lf._simulation_random_walk_state = pdf.attrs["_simulation_random_walk_state"]

            # Preserve scheduled event state
            if hasattr(pdf, "attrs") and "_simulation_scheduled_event_state" in pdf.attrs:
                lf._simulation_scheduled_event_state = pdf.attrs[
                    "_simulation_scheduled_event_state"
                ]

            ctx.info(
                "Simulation read completed (via Pandas, converted to Polars LazyFrame)",
                row_count=len(pdf),
            )
            return lf

        # Get full path
        if path:
            if connection:
                full_path = connection.get_path(path)
            else:
                full_path = path
        elif table:
            if connection:
                full_path = connection.get_path(table)
            else:
                raise ValueError(
                    f"Cannot read table '{table}': connection is required when using 'table' parameter. "
                    "Provide a valid connection object or use 'path' for file-based reads."
                )
        else:
            raise ValueError(
                "Read operation failed: neither 'path' nor 'table' was provided. "
                "Specify a file path or table name in your configuration."
            )

        # Handle glob patterns/lists
        # Polars scan methods often support glob strings directly.

        try:
            if format == "csv":
                # scan_csv supports glob patterns
                return pl.scan_csv(full_path, **options)

            elif format == "parquet":
                return pl.scan_parquet(full_path, **options)

            elif format == "json":
                # scan_ndjson for newline delimited json, read_json for standard
                # Assuming ndjson/jsonl for big data usually
                if options.get("json_lines", True):  # Default to ndjson scan
                    return pl.scan_ndjson(full_path, **options)
                else:
                    # Standard JSON doesn't support lazy scan well in all versions, fallback to read
                    return pl.read_json(full_path, **options).lazy()

            elif format == "delta":
                # scan_delta requires 'deltalake' extra usually or feature
                storage_options = options.get("storage_options", None)
                version = options.get("versionAsOf", None)

                # scan_delta is available in recent polars
                # It might accept storage_options in recent versions
                delta_opts = {}
                if storage_options:
                    delta_opts["storage_options"] = storage_options
                if version is not None:
                    delta_opts["version"] = version

                return pl.scan_delta(full_path, **delta_opts)

            elif format == "excel":
                # Excel format: delegate to Pandas (best Excel support), convert to Polars
                from odibi.engine.pandas_engine import PandasEngine
                from odibi.context import get_logging_context

                ctx = get_logging_context().with_context(engine="polars")
                ctx.debug("Reading Excel via Pandas engine (best Excel support)")

                # Get storage_options from connection for Azure/cloud authentication
                excel_storage_options = None
                if hasattr(connection, "pandas_storage_options"):
                    excel_storage_options = connection.pandas_storage_options()

                # Use Pandas engine for Excel reading
                pandas_engine = PandasEngine()
                pdf = pandas_engine._read_excel_with_patterns(
                    full_path,
                    sheet_pattern=options.pop("sheet_pattern", None),
                    sheet_pattern_case_sensitive=options.pop("sheet_pattern_case_sensitive", False),
                    add_source_file=options.pop("add_source_file", False),
                    is_glob="*" in str(full_path) or "?" in str(full_path),
                    ctx=ctx,
                    storage_options=excel_storage_options,
                    **options,
                )

                # Convert Pandas DataFrame to Polars LazyFrame
                ctx.info(f"Excel read completed (via Pandas): {path}", row_count=len(pdf))
                return pl.from_pandas(pdf).lazy()

            elif format == "api":
                # API format: delegate to Pandas (uses ApiFetcher), convert to Polars
                from odibi.connections.api_fetcher import create_api_fetcher
                from odibi.connections.http import HttpConnection
                from odibi.context import get_logging_context

                ctx = get_logging_context().with_context(engine="polars")
                ctx.debug("Reading API via ApiFetcher", endpoint=str(path))

                if not isinstance(connection, HttpConnection):
                    raise ValueError(
                        f"Cannot read API data: connection type '{type(connection).__name__}' "
                        "is not an HttpConnection. Use an HTTP connection for API format."
                    )

                # Extract API-specific options
                api_options = options.copy()
                params = api_options.pop("params", {})
                max_records = api_options.pop("max_records", None)
                method = api_options.pop("method", "GET")
                request_body = api_options.pop("request_body", None)

                # Create fetcher with connection's base_url and headers
                fetcher = create_api_fetcher(
                    base_url=connection.base_url,
                    headers=connection.headers,
                    options=api_options,
                )

                # Fetch data as Pandas DataFrame
                pdf = fetcher.fetch_dataframe(
                    endpoint=str(full_path),
                    params=params,
                    max_records=max_records,
                    method=method,
                    request_body=request_body,
                )

                ctx.info(f"API read completed: {path}", row_count=len(pdf))
                return pl.from_pandas(pdf).lazy()

            else:
                raise ValueError(
                    f"Unsupported format for Polars engine: '{format}'. "
                    "Supported formats: csv, parquet, json, delta, excel, api, sql, sql_server, azure_sql."
                )

        except Exception as e:
            raise ValueError(
                f"Failed to read {format} from '{full_path}': {e}. "
                "Check that the file exists, the format is correct, and you have read permissions."
            )

    def write(
        self,
        df: Any,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        mode: str = "overwrite",
        options: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Write data using Polars.

        Args:
            df: DataFrame or LazyFrame to write.
            connection: Connection object providing base path and storage options.
            format: Output format (csv, parquet, delta, json, sql, sql_server, azure_sql).
            table: Table name (mutually exclusive with path).
            path: File path (mutually exclusive with table).
            mode: Write mode (overwrite, append, upsert, append_once).
            options: Additional write options to pass to Polars writers.
            streaming_config: Streaming configuration (not used in Polars engine).

        Returns:
            Optional dict with write statistics for Delta writes, None otherwise.
        """
        options = options or {}

        if is_sql_format(format):
            return self._write_sql(df, connection, table, mode, options)

        if path:
            if connection:
                full_path = connection.get_path(path)
            else:
                full_path = path
        elif table:
            if connection:
                full_path = connection.get_path(table)
            else:
                raise ValueError(
                    f"Cannot write to table '{table}': connection is required when using 'table' parameter. "
                    "Provide a valid connection object or use 'path' for file-based writes."
                )
        else:
            raise ValueError(
                "Write operation failed: neither 'path' nor 'table' was provided. "
                "Specify a file path or table name in your configuration."
            )

        is_lazy = isinstance(df, pl.LazyFrame)

        # Handle upsert/append_once for non-Delta formats
        if mode in ["upsert", "append_once"] and format != "delta":
            df, mode = self._handle_generic_upsert(df, full_path, format, mode, options)
            is_lazy = isinstance(df, pl.LazyFrame)
            options = {k: v for k, v in options.items() if k != "keys"}

        # Handle plain append for file formats — Polars file writers overwrite by default,
        # so without read-existing + concat an append would silently lose prior data.
        if mode == "append" and format in ("parquet", "csv", "json"):
            df, mode = self._handle_file_append(df, full_path, format)
            is_lazy = isinstance(df, pl.LazyFrame)

        parent_dir = os.path.dirname(full_path)
        if parent_dir:
            os.makedirs(parent_dir, exist_ok=True)

        if format == "parquet":
            if is_lazy:
                df.sink_parquet(full_path, **options)
            else:
                df.write_parquet(full_path, **options)

        elif format == "csv":
            if is_lazy:
                df.sink_csv(full_path, **options)
            else:
                df.write_csv(full_path, **options)

        elif format == "json":
            if is_lazy:
                df.sink_ndjson(full_path, **options)
            else:
                df.write_ndjson(full_path, **options)

        elif format == "delta":
            if is_lazy:
                df = df.collect()

            storage_options = options.get("storage_options", None)
            delta_write_options = options.copy()
            if "storage_options" in delta_write_options:
                del delta_write_options["storage_options"]

            df.write_delta(
                full_path, mode=mode, storage_options=storage_options, **delta_write_options
            )

        else:
            raise ValueError(
                f"Unsupported write format for Polars engine: '{format}'. "
                "Supported formats: csv, parquet, json, delta."
            )

        return None

    def _write_sql(
        self,
        df: Any,
        connection: Any,
        table: Optional[str],
        mode: str,
        options: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """Handle SQL writing including merge and enhanced overwrite for Polars (Phase 4)."""
        from odibi.utils.logging_context import get_logging_context

        ctx = get_logging_context().with_context(engine="polars")

        if not hasattr(connection, "write_table"):
            raise ValueError(
                f"Connection type '{type(connection).__name__}' does not support SQL operations"
            )

        if not table:
            raise ValueError(
                "SQL write operation failed: 'table' parameter is required but was not provided. "
                "Specify the target table name in your configuration."
            )

        if mode == "merge":
            if getattr(connection, "sql_dialect", "mssql") != "mssql":
                raise NotImplementedError(
                    f"SQL MERGE mode is currently only supported for SQL Server connections. "
                    f"Connection dialect '{getattr(connection, 'sql_dialect', 'unknown')}' "
                    f"does not support MERGE. Use mode='append' or mode='overwrite' instead."
                )

            merge_keys = options.get("merge_keys")
            merge_options = options.get("merge_options")

            if not merge_keys:
                raise ValueError(
                    "MERGE mode requires 'merge_keys' in options. "
                    "Specify the key columns for the MERGE ON clause."
                )

            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            writer = SqlServerMergeWriter(connection)
            ctx.debug(
                "Executing SQL Server MERGE (Polars)",
                target=table,
                merge_keys=merge_keys,
            )

            result = writer.merge_polars(
                df=df,
                target_table=table,
                merge_keys=merge_keys,
                options=merge_options,
            )

            ctx.info(
                "SQL Server MERGE completed (Polars)",
                target=table,
                inserted=result.inserted,
                updated=result.updated,
                deleted=result.deleted,
            )

            return {
                "mode": "merge",
                "inserted": result.inserted,
                "updated": result.updated,
                "deleted": result.deleted,
                "total_affected": result.total_affected,
            }

        if mode == "overwrite" and options.get("overwrite_options"):
            if getattr(connection, "sql_dialect", "mssql") != "mssql":
                raise NotImplementedError(
                    f"Enhanced overwrite strategies are currently only supported for SQL Server. "
                    f"Use standard mode='overwrite' without overwrite_options for "
                    f"'{getattr(connection, 'sql_dialect', 'unknown')}' connections."
                )
            from odibi.writers.sql_server_writer import SqlServerMergeWriter

            overwrite_options = options.get("overwrite_options")
            writer = SqlServerMergeWriter(connection)

            ctx.debug(
                "Executing SQL Server enhanced overwrite (Polars)",
                target=table,
                strategy=(
                    overwrite_options.strategy.value
                    if hasattr(overwrite_options, "strategy")
                    else "truncate_insert"
                ),
            )

            result = writer.overwrite_polars(
                df=df,
                target_table=table,
                options=overwrite_options,
            )

            ctx.info(
                "SQL Server enhanced overwrite completed (Polars)",
                target=table,
                strategy=result.strategy,
                rows_written=result.rows_written,
            )

            return {
                "mode": "overwrite",
                "strategy": result.strategy,
                "rows_written": result.rows_written,
            }

        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        default_schema = getattr(connection, "default_schema", "dbo")
        if "." in table:
            schema, table_name = table.split(".", 1)
        else:
            schema, table_name = default_schema, table

        if_exists = "replace"
        if mode == "append":
            if_exists = "append"
        elif mode == "fail":
            if_exists = "fail"

        df_pandas = df.to_pandas()
        chunksize = options.get("chunksize", 1000)

        connection.write_table(
            df=df_pandas,
            table_name=table_name,
            schema=schema,
            if_exists=if_exists,
            chunksize=chunksize,
        )
        return None

    def _handle_generic_upsert(
        self,
        df: Any,
        full_path: str,
        format: str,
        mode: str,
        options: Dict[str, Any],
    ) -> tuple:
        """Handle upsert/append_once for standard files by merging with existing data."""
        if "keys" not in options:
            raise ValueError(f"Mode '{mode}' requires 'keys' list in options")

        keys = options["keys"]
        if isinstance(keys, str):
            keys = [keys]

        # Materialize LazyFrame for join operations
        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        # Try to read existing file
        existing_df = None
        try:
            if format == "csv":
                existing_df = pl.read_csv(full_path)
            elif format == "parquet":
                existing_df = pl.read_parquet(full_path)
            elif format == "json":
                existing_df = pl.read_ndjson(full_path)
        except Exception:
            return df, "overwrite"

        if existing_df is None:
            return df, "overwrite"

        if mode == "append_once":
            missing_keys = set(keys) - set(df.columns)
            if missing_keys:
                raise KeyError(f"Keys {missing_keys} not found in input data")

            new_rows = df.join(existing_df.select(keys), on=keys, how="anti")
            # diagonal_relaxed aligns by column NAME (not position) so a differing
            # column order between the existing file and new data doesn't raise.
            return pl.concat([existing_df, new_rows], how="diagonal_relaxed"), "overwrite"

        elif mode == "upsert":
            missing_keys = set(keys) - set(df.columns)
            if missing_keys:
                raise KeyError(f"Keys {missing_keys} not found in input data")

            rows_to_keep = existing_df.join(df.select(keys), on=keys, how="anti")
            return pl.concat([rows_to_keep, df], how="diagonal_relaxed"), "overwrite"

        return df, mode

    def _handle_file_append(self, df: Any, full_path: str, format: str) -> tuple:
        """Implement append for file formats by read-existing + concat.

        Polars' write_parquet/write_csv/write_ndjson have no append mode and overwrite
        the target, so a plain mode="append" would silently destroy existing data.
        Returns (combined_df, "overwrite") — or the original df if there's nothing to
        append to yet.
        """
        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        existing_df = None
        try:
            if format == "csv":
                existing_df = pl.read_csv(full_path)
            elif format == "parquet":
                existing_df = pl.read_parquet(full_path)
            elif format == "json":
                existing_df = pl.read_ndjson(full_path)
        except Exception:
            return df, "overwrite"

        if existing_df is None:
            return df, "overwrite"

        return pl.concat([existing_df, df], how="diagonal_relaxed"), "overwrite"

    def execute_sql(self, sql: str, context: Context) -> Any:
        """Execute SQL query using Polars SQLContext.

        Args:
            sql: SQL query string
            context: Execution context with registered DataFrames

        Returns:
            pl.LazyFrame
        """
        ctx = pl.SQLContext()

        # Register datasets from context
        # We iterate over all registered names in the context
        try:
            names = context.list_names()
            for name in names:
                df = context.get(name)
                # Register LazyFrame or DataFrame
                # Polars SQLContext supports registering LazyFrame, DataFrame, and some others
                # We might need to convert if it's not a Polars object, but we assume Polars engine uses Polars objects
                ctx.register(name, df)
        except Exception:
            # If context doesn't support listing or getting, we proceed with empty context
            # (e.g. if context is not fully compatible or empty)
            pass

        return ctx.execute(sql, eager=False)

    def execute_operation(self, operation: str, params: Dict[str, Any], df: Any) -> Any:
        """Execute built-in operation.

        Args:
            operation: Name of the operation to execute (e.g., 'pivot', 'unpivot', 'explode').
            params: Dictionary of parameters specific to the operation.
            df: DataFrame or LazyFrame to operate on.

        Returns:
            Transformed DataFrame or LazyFrame.
        """
        # Ensure LazyFrame for consistency if possible, but operations work on both usually.
        # If DataFrame, some operations might need different methods.

        if operation == "pivot":
            # Pivot requires materialization usually in other engines, but Polars LazyFrame has 'collect' or similar constraints?
            # Polars lazy pivot is not fully supported in older versions without collect, but check recent.
            # Pivot changes shape drastically.
            # params: pivot_column, value_column, group_by, agg_func

            # If lazy, we might need to collect for pivot if lazy pivot isn't supported or experimental.
            # But let's try to keep it lazy if possible.
            # As of recent Polars, pivot is available on DataFrame, experimental on LazyFrame?
            # Actually, 'unstack' or 'pivot' on LazyFrame is limited.
            # Safe bet: materialize if needed, or use lazy pivot if available.

            # Let's collect if input is lazy, because pivot usually implies strict schema change hard to predict.
            if isinstance(df, pl.LazyFrame):
                df = df.collect()

            return df.pivot(
                index=params.get("group_by"),
                on=params["pivot_column"],
                values=params["value_column"],
                aggregate_function=params.get("agg_func", "first"),
            )  # Returns DataFrame

        elif operation == "drop_duplicates":
            subset = params.get("subset")
            if isinstance(df, pl.LazyFrame):
                return df.unique(subset=subset)
            return df.unique(subset=subset)

        elif operation == "fillna":
            value = params.get("value")
            # Polars uses fill_null
            if isinstance(value, dict):
                # Fill specific columns
                # value = {'col1': 0, 'col2': 'unknown'}
                # We need to chain with_columns
                exprs = []
                for col, val in value.items():
                    exprs.append(pl.col(col).fill_null(val))
                return df.with_columns(exprs)
            else:
                # Fill all columns? Polars fill_null requires specifying columns or using all()
                return df.fill_null(value)

        elif operation == "drop":
            columns = params.get("columns") or params.get("labels")
            return df.drop(columns)

        elif operation == "rename":
            columns = params.get("columns") or params.get("mapper")
            return df.rename(columns)

        elif operation == "sort":
            by = params.get("by")
            descending = not params.get("ascending", True)
            if isinstance(df, pl.LazyFrame):
                return df.sort(by, descending=descending)
            return df.sort(by, descending=descending)

        elif operation == "sample":
            # Sample n or frac
            n = params.get("n")
            frac = params.get("frac")
            seed = params.get("random_state")

            # Lazy sample supported
            if n is not None:
                # Note: Polars Lazy sample might be approximate or require 'collect' depending on version/backend?
                # But usually supported.
                if isinstance(df, pl.LazyFrame):
                    # LazyFrame.sample takes n (int) or fraction.
                    # But polars 0.19+ changed sample signature?
                    # It's generally `sample(n=..., fraction=..., seed=...)`
                    return (
                        df.collect().sample(n=n, seed=seed).lazy()
                    )  # Collecting for exact sample n on lazy might be needed if not supported?
                    # Actually, fetch(n) is head. Sample is random.
                    # Let's materialize for safety with sample as it's often for checks.
                    pass
                return df.sample(n=n, seed=seed)
            elif frac is not None:
                if isinstance(df, pl.LazyFrame):
                    # Lazy sampling by fraction is supported
                    pass  # fall through
                return df.sample(fraction=frac, seed=seed)

        elif operation == "filter":
            # Legacy or simple filter
            pass

        else:
            # Fallback: check if operation is a registered transformer
            from odibi.context import EngineContext, PandasContext
            from odibi.registry import FunctionRegistry

            if FunctionRegistry.has_function(operation):
                func = FunctionRegistry.get_function(operation)
                param_model = FunctionRegistry.get_param_model(operation)

                # Create EngineContext from current df (use PandasContext as placeholder)
                engine_ctx = EngineContext(
                    context=PandasContext(),
                    df=df,
                    engine=self,
                    engine_type=self.engine_type,
                )

                # Validate and instantiate params
                if param_model:
                    validated_params = param_model(**params)
                    result_ctx = func(engine_ctx, validated_params)
                else:
                    result_ctx = func(engine_ctx, **params)

                return result_ctx.df

        return df

    def get_schema(self, df: Any) -> Any:
        """Get DataFrame schema.

        Args:
            df: DataFrame or LazyFrame to get schema from.

        Returns:
            Dictionary mapping column names to data type strings.
        """
        # Polars schema is a dict {name: DataType}
        # We can return a dict of strings for compatibility
        schema = df.collect_schema() if isinstance(df, pl.LazyFrame) else df.schema
        return {name: str(dtype) for name, dtype in schema.items()}

    def get_shape(self, df: Any) -> tuple:
        """Get DataFrame shape.

        Args:
            df: DataFrame or LazyFrame to get shape from.

        Returns:
            Tuple of (rows, columns) as integers.
        """
        if isinstance(df, pl.LazyFrame):
            # Expensive to count rows in LazyFrame without scan
            # But usually shape implies (rows, cols)
            # columns is cheap. rows requires partial scan or metadata.
            # Fetching 1 row might give columns.
            # For exact row count, we need collect(count)
            cols = len(df.collect_schema().names())
            rows = df.select(pl.len()).collect().item()
            return (rows, cols)
        return df.shape

    def count_rows(self, df: Any) -> int:
        """Count rows in DataFrame.

        Args:
            df: DataFrame or LazyFrame to count rows from.

        Returns:
            Number of rows as integer.
        """
        if isinstance(df, pl.LazyFrame):
            return df.select(pl.len()).collect().item()
        return len(df)

    def count_nulls(self, df: Any, columns: List[str]) -> Dict[str, int]:
        """Count nulls in specified columns.

        Args:
            df: DataFrame or LazyFrame to analyze.
            columns: List of column names to count nulls for.

        Returns:
            Dictionary mapping column names to null counts.
        """
        if isinstance(df, pl.LazyFrame):
            # efficient null count
            return df.select([pl.col(c).null_count() for c in columns]).collect().to_dicts()[0]

        return df.select([pl.col(c).null_count() for c in columns]).to_dicts()[0]

    def validate_schema(self, df: Any, schema_rules: Dict[str, Any]) -> List[str]:
        """Validate DataFrame schema.

        Args:
            df: DataFrame or LazyFrame to validate.
            schema_rules: Dictionary containing schema validation rules
                (e.g., required_columns, expected_types, disallowed_columns).

        Returns:
            List of validation failure messages (empty if all validations pass).
        """
        failures = []

        # Schema is dict-like in Polars
        current_schema = df.collect_schema() if isinstance(df, pl.LazyFrame) else df.schema
        current_cols = current_schema.keys()

        if "required_columns" in schema_rules:
            required = schema_rules["required_columns"]
            missing = set(required) - set(current_cols)
            if missing:
                failures.append(f"Missing required columns: {', '.join(missing)}")

        if "types" in schema_rules:
            for col, expected_type in schema_rules["types"].items():
                if col not in current_cols:
                    failures.append(f"Column '{col}' not found for type validation")
                    continue

                actual_type = str(current_schema[col])
                # Basic type check - simplistic string matching
                if expected_type.lower() not in actual_type.lower():
                    failures.append(
                        f"Column '{col}' has type '{actual_type}', expected '{expected_type}'"
                    )

        return failures

    def validate_data(self, df: Any, validation_config: Any) -> List[str]:
        """Validate data against rules.

        Args:
            df: DataFrame or LazyFrame
            validation_config: ValidationConfig object

        Returns:
            List of validation failure messages
        """
        failures = []

        if isinstance(df, pl.LazyFrame):
            schema = df.collect_schema()
            columns = schema.names()
        else:
            columns = df.columns

        if getattr(validation_config, "not_empty", False):
            count = self.count_rows(df)
            if count == 0:
                failures.append("DataFrame is empty")

        if getattr(validation_config, "no_nulls", None):
            cols = validation_config.no_nulls
            null_counts = self.count_nulls(df, cols)
            for col, count in null_counts.items():
                if count > 0:
                    failures.append(f"Column '{col}' has {count} null values")

        if getattr(validation_config, "schema_validation", None):
            schema_failures = self.validate_schema(df, validation_config.schema_validation)
            failures.extend(schema_failures)

        if getattr(validation_config, "ranges", None):
            for col, bounds in validation_config.ranges.items():
                if col in columns:
                    min_val = bounds.get("min")
                    max_val = bounds.get("max")

                    if min_val is not None:
                        if isinstance(df, pl.LazyFrame):
                            min_violations = (
                                df.filter(pl.col(col) < min_val).select(pl.len()).collect().item()
                            )
                        else:
                            min_violations = len(df.filter(pl.col(col) < min_val))
                        if min_violations > 0:
                            failures.append(f"Column '{col}' has values < {min_val}")

                    if max_val is not None:
                        if isinstance(df, pl.LazyFrame):
                            max_violations = (
                                df.filter(pl.col(col) > max_val).select(pl.len()).collect().item()
                            )
                        else:
                            max_violations = len(df.filter(pl.col(col) > max_val))
                        if max_violations > 0:
                            failures.append(f"Column '{col}' has values > {max_val}")
                else:
                    failures.append(f"Column '{col}' not found for range validation")

        if getattr(validation_config, "allowed_values", None):
            for col, allowed in validation_config.allowed_values.items():
                if col in columns:
                    if isinstance(df, pl.LazyFrame):
                        invalid_count = (
                            df.filter(~pl.col(col).is_in(allowed)).select(pl.len()).collect().item()
                        )
                    else:
                        invalid_count = len(df.filter(~pl.col(col).is_in(allowed)))
                    if invalid_count > 0:
                        failures.append(f"Column '{col}' has invalid values")
                else:
                    failures.append(f"Column '{col}' not found for allowed values validation")

        return failures

    def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
        """Get sample rows as list of dictionaries.

        Args:
            df: DataFrame or LazyFrame to sample from.
            n: Number of rows to sample (default: 10).

        Returns:
            List of dictionaries, each representing a row.
        """
        if isinstance(df, pl.LazyFrame):
            return df.limit(n).collect().to_dicts()
        return df.head(n).to_dicts()

    def profile_nulls(self, df: Any) -> Dict[str, float]:
        """Calculate null percentage for each column.

        Args:
            df: DataFrame or LazyFrame to profile.

        Returns:
            Dictionary mapping column names to null percentages (0.0 to 1.0).
        """
        if isinstance(df, pl.LazyFrame):
            # null_count() / count()
            # We can do this in one expression
            total_count = df.select(pl.len()).collect().item()
            if total_count == 0:
                return {col: 0.0 for col in df.collect_schema().names()}

            cols = df.collect_schema().names()
            null_counts = df.select([pl.col(c).null_count().alias(c) for c in cols]).collect()
            return {col: null_counts[col][0] / total_count for col in cols}

        total_count = len(df)
        if total_count == 0:
            return {col: 0.0 for col in df.columns}

        null_counts = df.null_count()
        return {col: null_counts[col][0] / total_count for col in df.columns}

    def table_exists(
        self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
    ) -> bool:
        """Check if table or location exists.

        Args:
            connection: Connection object to resolve paths.
            table: Table name to check (mutually exclusive with path).
            path: File path to check (mutually exclusive with table).

        Returns:
            True if the table or path exists, False otherwise.
        """
        if path:
            full_path = connection.get_path(path)
            return os.path.exists(full_path)
        return False

    def harmonize_schema(self, df: Any, target_schema: Dict[str, str], policy: Any) -> Any:
        """Harmonize DataFrame schema.

        Args:
            df: DataFrame or LazyFrame to harmonize.
            target_schema: Dictionary mapping column names to target data types.
            policy: SchemaPolicyConfig object defining harmonization behavior
                (mode, on_missing_columns, on_new_columns).

        Returns:
            DataFrame or LazyFrame with harmonized schema.
        """
        # policy: SchemaPolicyConfig
        from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

        # Helper to get current columns/schema
        if isinstance(df, pl.LazyFrame):
            current_schema = df.collect_schema()
        else:
            current_schema = df.schema

        current_cols = current_schema.names()
        target_cols = list(target_schema.keys())

        missing = set(target_cols) - set(current_cols)
        new_cols = set(current_cols) - set(target_cols)

        # 1. Validation
        if missing and getattr(policy, "on_missing_columns", None) == OnMissingColumns.FAIL:
            raise ValueError(
                f"Schema Policy Violation: DataFrame is missing required columns {missing}. "
                f"Available columns: {current_cols}. Add missing columns or set on_missing_columns policy."
            )

        if new_cols and getattr(policy, "on_new_columns", None) == OnNewColumns.FAIL:
            raise ValueError(
                f"Schema Policy Violation: DataFrame contains unexpected columns {new_cols}. "
                f"Expected columns: {target_cols}. Remove extra columns or set on_new_columns policy."
            )

        # 2. Transformations
        exprs = []

        # Handle Missing (Add nulls)
        # Evolve means we keep new columns, Enforce means we select only target
        mode = getattr(policy, "mode", SchemaMode.ENFORCE)

        if (
            mode == SchemaMode.EVOLVE
            and getattr(policy, "on_new_columns", None) == OnNewColumns.ADD_NULLABLE
        ):
            # Add missing (if missing cols exist, we fill them with nulls)
            # on_missing_columns controls what to do with missing target cols.
            # If mode is EVOLVE, we typically keep everything?
            # But harmonize_schema is about matching a TARGET schema.
            # If target has cols that df doesn't:
            # If on_missing_columns == FILL_NULL -> Add them as null.
            pass

        # We should respect on_missing_columns regardless of mode?
        if missing and getattr(policy, "on_missing_columns", None) == OnMissingColumns.FILL_NULL:
            for col in missing:
                exprs.append(pl.lit(None).alias(col))

        if exprs:
            df = df.with_columns(exprs)

        # Now Select
        if mode == SchemaMode.ENFORCE:
            # Select only target columns.
            # Missing columns were added above if configured.
            # New columns (not in target) are dropped implicitly by selecting target_cols.
            # But wait, we added exprs to df (lazy).

            final_cols = []
            for col in target_cols:
                final_cols.append(pl.col(col))

            df = df.select(final_cols)

        elif mode == SchemaMode.EVOLVE:
            # We keep new columns.
            # If target has columns that were missing in df, we added them above (if FILL_NULL).
            # If df has columns not in target (new_cols), we keep them.
            pass

        return df

    def anonymize(
        self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
    ) -> Any:
        """Anonymize specified columns.

        Args:
            df: DataFrame or LazyFrame to anonymize.
            columns: List of column names to anonymize.
            method: Anonymization method ('mask', 'hash', 'redact').
            salt: Optional salt string for hash-based anonymization.

        Returns:
            DataFrame or LazyFrame with anonymized columns.
        """
        if method == "mask":
            # Mask all but last 4 characters: '******1234'
            # Regex look-around not supported in some envs.
            # Manual approach:
            # If len > 4: repeat('*', len-4) + suffix(4)
            # Else: keep original (or mask all? Pandas engine masked all but last 4, which implies keeping small strings?)
            # Pandas: .str.replace(r".(?=.{4})", "*") -> replaces chars that are followed by 4 chars.
            # If str is "123", no char is followed by 4 chars -> "123".
            # If str is "12345", '1' is followed by '2345' (4 chars) -> "*2345".

            return df.with_columns(
                [
                    pl.when(pl.col(c).cast(pl.Utf8).str.len_chars() > 4)
                    .then(
                        pl.concat_str(
                            [
                                # Cast to signed Int64 BEFORE subtracting, then clip to 0:
                                # len_chars() is UInt32, so `len - 4` for short strings would
                                # underflow to a huge unsigned value. Polars evaluates .then()
                                # for ALL rows (not just len>4), so without this repeat_by would
                                # try to allocate tens of GB and abort the process.
                                pl.lit("*")
                                .repeat_by(
                                    (pl.col(c).str.len_chars().cast(pl.Int64) - 4).clip(
                                        lower_bound=0
                                    )
                                )
                                .list.join(""),
                                pl.col(c).str.slice(-4),
                            ]
                        )
                    )
                    .otherwise(pl.col(c).cast(pl.Utf8))
                    .alias(c)
                    for c in columns
                ]
            )

        elif method == "hash":
            # Polars hash() is non-cryptographic usually (xxHash).
            # For cryptographic hash (sha256), we might need map_elements (slow) or plugin.
            # Requirement is just 'hash', often consistent for analytics.
            # Gap Analysis mentions "salt".
            # PandasEngine used sha256 with salt.
            # Polars `hash` is fast 64-bit hash.
            # If we need SHA256, we must use map_elements (python UDF) or custom.
            # For "High Performance", map_elements is bad.
            # However, without native plugin, we have no choice for SHA256.
            # Let's implement SHA256 via map_elements for compatibility,
            # OR use Polars internal hash if user accepts non-crypto.
            # But "salt" implies security/crypto usage.

            def _hash_val(val):
                if val is None:
                    return None
                to_hash = str(val)
                if salt:
                    to_hash += salt
                return hashlib.sha256(to_hash.encode("utf-8")).hexdigest()

            # Apply to each column. Warning: Slow path.
            # But Polars UDFs are still faster than Pandas apply often due to no GIL? No, Python UDF has GIL.
            return df.with_columns(
                [pl.col(c).map_elements(_hash_val, return_dtype=pl.Utf8).alias(c) for c in columns]
            )

        elif method == "redact":
            return df.with_columns([pl.lit("[REDACTED]").alias(c) for c in columns])

        return df

    def get_table_schema(
        self,
        connection: Any,
        table: Optional[str] = None,
        path: Optional[str] = None,
        format: Optional[str] = None,
    ) -> Optional[Dict[str, str]]:
        """Get schema of an existing table/file.

        Args:
            connection: Connection object
            table: Table name
            path: File path
            format: Data format (optional, helps with file-based sources)

        Returns:
            Schema dict or None if table doesn't exist or schema fetch fails.
        """
        from odibi.utils.logging_context import get_logging_context

        ctx = get_logging_context().with_context(engine="polars")

        try:
            if table and is_sql_format(format):
                query = connection.build_select_query(table, limit=0)
                df = connection.read_sql(query)
                return {col: str(dtype) for col, dtype in zip(df.columns, df.dtypes)}

            if path:
                full_path = connection.get_path(path) if connection else path
                if not os.path.exists(full_path):
                    return None

                if format == "delta":
                    try:
                        from deltalake import DeltaTable

                        dt = DeltaTable(full_path)
                        arrow_schema = dt.schema().to_pyarrow()
                        return {field.name: str(field.type) for field in arrow_schema}
                    except ImportError:
                        ctx.warning(
                            "deltalake library not installed for schema introspection",
                            path=full_path,
                        )
                        return None

                elif format == "parquet":
                    try:
                        import pyarrow.parquet as pq
                        import glob as glob_mod

                        target_path = full_path
                        if os.path.isdir(full_path):
                            files = glob_mod.glob(os.path.join(full_path, "*.parquet"))
                            if not files:
                                return None
                            target_path = files[0]

                        schema = pq.read_schema(target_path)
                        return {field.name: str(field.type) for field in schema}
                    except ImportError:
                        lf = pl.scan_parquet(full_path)
                        schema = lf.collect_schema()
                        return {name: str(dtype) for name, dtype in schema.items()}

                elif format == "csv":
                    lf = pl.scan_csv(full_path)
                    schema = lf.collect_schema()
                    return {name: str(dtype) for name, dtype in schema.items()}

        except (FileNotFoundError, PermissionError):
            return None
        except Exception as e:
            ctx.warning(f"Failed to infer schema for {table or path}: {e}")
            return None

        return None

    def maintain_table(
        self,
        connection: Any,
        format: str,
        table: Optional[str] = None,
        path: Optional[str] = None,
        config: Optional[Any] = None,
    ) -> None:
        """Run table maintenance operations (optimize, vacuum) for Delta tables.

        Args:
            connection: Connection object
            format: Table format
            table: Table name
            path: Table path
            config: AutoOptimizeConfig object

        Returns:
            None
        """
        from odibi.utils.logging_context import get_logging_context

        ctx = get_logging_context().with_context(engine="polars")

        if format != "delta" or not config or not getattr(config, "enabled", False):
            return

        if not path and not table:
            return

        full_path = connection.get_path(path if path else table) if connection else (path or table)

        ctx.info("Starting table maintenance", path=str(full_path))

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.warning(
                "Auto-optimize skipped: 'deltalake' library not installed",
                path=str(full_path),
            )
            return

        try:
            import time

            start = time.time()

            storage_opts = {}
            if hasattr(connection, "pandas_storage_options"):
                storage_opts = connection.pandas_storage_options()

            dt = DeltaTable(full_path, storage_options=storage_opts)

            ctx.info("Running Delta OPTIMIZE (compaction)", path=str(full_path))
            dt.optimize.compact()

            retention = getattr(config, "vacuum_retention_hours", None)
            if retention is not None and retention > 0:
                ctx.info(
                    "Running Delta VACUUM",
                    path=str(full_path),
                    retention_hours=retention,
                )
                dt.vacuum(
                    retention_hours=retention,
                    enforce_retention_duration=True,
                    dry_run=False,
                )

            elapsed = (time.time() - start) * 1000
            ctx.info(
                "Table maintenance completed",
                path=str(full_path),
                elapsed_ms=round(elapsed, 2),
            )

        except Exception as e:
            ctx.warning(
                "Auto-optimize failed",
                path=str(full_path),
                error=str(e),
            )

    def get_source_files(self, df: Any) -> List[str]:
        """Get list of source files that generated this DataFrame.

        For Polars, this checks if source file info was stored
        in the DataFrame's metadata during read.

        Args:
            df: DataFrame or LazyFrame

        Returns:
            List of file paths (or empty list if not applicable/supported)
        """
        if isinstance(df, pl.LazyFrame):
            return []

        if hasattr(df, "attrs"):
            return df.attrs.get("odibi_source_files", [])

        return []

    def vacuum_delta(
        self,
        connection: Any,
        path: str,
        retention_hours: int = 168,
        dry_run: bool = False,
        enforce_retention_duration: bool = True,
    ) -> Dict[str, Any]:
        """VACUUM a Delta table to remove old files.

        Args:
            connection: Connection object
            path: Delta table path
            retention_hours: Retention period (default 168 = 7 days)
            dry_run: If True, only show files to be deleted
            enforce_retention_duration: If False, allows retention < 168 hours (testing only)

        Returns:
            Dictionary with files_deleted count
        """
        from odibi.utils.logging_context import get_logging_context
        import time

        ctx = get_logging_context().with_context(engine="polars")
        start = time.time()

        ctx.debug(
            "Starting Delta VACUUM",
            path=path,
            retention_hours=retention_hours,
            dry_run=dry_run,
        )

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[polars]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path) if connection else path

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        deleted_files = dt.vacuum(
            retention_hours=retention_hours,
            dry_run=dry_run,
            enforce_retention_duration=enforce_retention_duration,
        )

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta VACUUM completed",
            path=str(full_path),
            files_deleted=len(deleted_files),
            dry_run=dry_run,
            elapsed_ms=round(elapsed, 2),
        )

        return {"files_deleted": len(deleted_files)}

    def get_delta_history(
        self, connection: Any, path: str, limit: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Get Delta table history.

        Args:
            connection: Connection object
            path: Delta table path
            limit: Maximum number of versions to return

        Returns:
            List of version metadata dictionaries
        """
        from odibi.utils.logging_context import get_logging_context
        import time

        ctx = get_logging_context().with_context(engine="polars")
        start = time.time()

        ctx.debug("Getting Delta table history", path=path, limit=limit)

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[polars]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path) if connection else path

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        history = dt.history(limit=limit)

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta history retrieved",
            path=str(full_path),
            versions_returned=len(history) if history else 0,
            elapsed_ms=round(elapsed, 2),
        )

        return history

    def add_write_metadata(
        self,
        df: "pl.DataFrame",
        metadata_config: Any,
        source_connection: Optional[str] = None,
        source_table: Optional[str] = None,
        source_path: Optional[str] = None,
        is_file_source: bool = False,
    ) -> "pl.DataFrame":
        """Add metadata columns to DataFrame before writing (Bronze layer lineage).

        Args:
            df: Polars DataFrame or LazyFrame
            metadata_config: WriteMetadataConfig or True (for all defaults)
            source_connection: Name of the source connection
            source_table: Name of the source table (SQL sources)
            source_path: Path of the source file (file sources)
            is_file_source: True if source is a file-based read

        Returns:
            DataFrame with metadata columns added
        """
        from odibi.config import WriteMetadataConfig

        if metadata_config is True:
            config = WriteMetadataConfig()
        elif isinstance(metadata_config, WriteMetadataConfig):
            config = metadata_config
        else:
            return df

        columns = []
        if config.extracted_at:
            columns.append(pl.lit(datetime.now(timezone.utc)).alias("_extracted_at"))
        if config.source_file and is_file_source and source_path:
            columns.append(pl.lit(source_path).alias("_source_file"))
        if config.source_connection and source_connection:
            columns.append(pl.lit(source_connection).alias("_source_connection"))
        if config.source_table and source_table:
            columns.append(pl.lit(source_table).alias("_source_table"))

        if columns:
            df = df.with_columns(columns)

        return df

    def restore_delta(self, connection: Any, path: str, version: int) -> None:
        """Restore Delta table to a specific version.

        Args:
            connection: Connection object
            path: Delta table path
            version: Version number to restore to
        """
        from odibi.utils.logging_context import get_logging_context

        ctx = get_logging_context().with_context(engine="polars")
        start = time.time()

        ctx.info("Starting Delta table restore", path=path, target_version=version)

        try:
            from deltalake import DeltaTable
        except ImportError:
            ctx.error("Delta Lake library not installed", path=path)
            raise ImportError(
                "Delta Lake support requires 'pip install odibi[polars]' "
                "or 'pip install deltalake'. See README.md for installation instructions."
            )

        full_path = connection.get_path(path) if connection else path

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)
        dt.restore(version)

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Delta table restored",
            path=str(full_path),
            restored_to_version=version,
            elapsed_ms=round(elapsed, 2),
        )

    def filter_greater_than(self, df: "pl.DataFrame", column: str, value: Any) -> "pl.DataFrame":
        """Filter DataFrame where column > value.

        Automatically casts string columns to datetime for proper comparison.

        Args:
            df: Polars DataFrame or LazyFrame
            column: Column name to filter on
            value: Value to compare against

        Returns:
            Filtered DataFrame
        """
        is_lazy = isinstance(df, pl.LazyFrame)
        schema = df.collect_schema() if is_lazy else df.schema

        if column not in schema.names():
            raise ValueError(f"Column '{column}' not found in DataFrame")

        try:
            col_dtype = schema[column]

            if col_dtype == pl.Utf8:
                df = df.with_columns(pl.col(column).str.to_datetime(time_unit="us", strict=False))
                value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))
            elif col_dtype in (pl.Datetime, pl.Date) and isinstance(value, str):
                value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))

            return df.filter(pl.col(column) > value)
        except Exception as e:
            raise ValueError(f"Failed to filter {column} > {value}: {e}")

    def filter_coalesce(
        self,
        df: "pl.DataFrame",
        col1: str,
        col2: str,
        op: str,
        value: Any,
    ) -> "pl.DataFrame":
        """Filter using COALESCE(col1, col2) op value.

        Automatically casts string columns to datetime for proper comparison.

        Args:
            df: Polars DataFrame or LazyFrame
            col1: Primary column name
            col2: Fallback column name
            op: Comparison operator (>=, >, <=, <, ==, =)
            value: Value to compare against

        Returns:
            Filtered DataFrame
        """
        is_lazy = isinstance(df, pl.LazyFrame)
        schema = df.collect_schema() if is_lazy else df.schema

        if col1 not in schema.names():
            raise ValueError(f"Column '{col1}' not found")

        def _cast_if_string(col_name: str) -> pl.Expr:
            if schema[col_name] == pl.Utf8:
                return pl.col(col_name).str.to_datetime(strict=False)
            return pl.col(col_name)

        expr1 = _cast_if_string(col1)
        if col2 in schema.names():
            expr2 = _cast_if_string(col2)
            coalesced = pl.coalesce(expr1, expr2)
        else:
            coalesced = expr1

        try:
            # Determine if value needs datetime parsing by checking schema types
            col1_dtype = schema[col1]
            needs_datetime = False
            if col1_dtype == pl.Utf8:
                needs_datetime = True
            elif col1_dtype in (pl.Datetime, pl.Date):
                needs_datetime = True

            if needs_datetime and isinstance(value, str):
                value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))

            ops = {
                ">=": coalesced >= value,
                ">": coalesced > value,
                "<=": coalesced <= value,
                "<": coalesced < value,
                "==": coalesced == value,
                "=": coalesced == value,
            }

            if op not in ops:
                raise ValueError(f"Unsupported operator: {op}")

            return df.filter(ops[op])
        except Exception as e:
            raise ValueError(f"Failed to filter COALESCE({col1}, {col2}) {op} {value}: {e}")

__init__(connections=None, config=None)

Initialize Polars engine.

Parameters:

Name Type Description Default
connections Optional[Dict[str, Any]]

Dictionary of connection objects

None
config Optional[Dict[str, Any]]

Engine configuration (optional)

None
Source code in odibi/engine/polars_engine.py
def __init__(
    self,
    connections: Optional[Dict[str, Any]] = None,
    config: Optional[Dict[str, Any]] = None,
):
    """Initialize Polars engine.

    Args:
        connections: Dictionary of connection objects
        config: Engine configuration (optional)
    """
    if pl is None:
        raise ImportError("Polars not installed. Run 'pip install polars'.")

    self.connections = connections or {}
    self.config = config or {}

add_write_metadata(df, metadata_config, source_connection=None, source_table=None, source_path=None, is_file_source=False)

Add metadata columns to DataFrame before writing (Bronze layer lineage).

Parameters:

Name Type Description Default
df DataFrame

Polars DataFrame or LazyFrame

required
metadata_config Any

WriteMetadataConfig or True (for all defaults)

required
source_connection Optional[str]

Name of the source connection

None
source_table Optional[str]

Name of the source table (SQL sources)

None
source_path Optional[str]

Path of the source file (file sources)

None
is_file_source bool

True if source is a file-based read

False

Returns:

Type Description
DataFrame

DataFrame with metadata columns added

Source code in odibi/engine/polars_engine.py
def add_write_metadata(
    self,
    df: "pl.DataFrame",
    metadata_config: Any,
    source_connection: Optional[str] = None,
    source_table: Optional[str] = None,
    source_path: Optional[str] = None,
    is_file_source: bool = False,
) -> "pl.DataFrame":
    """Add metadata columns to DataFrame before writing (Bronze layer lineage).

    Args:
        df: Polars DataFrame or LazyFrame
        metadata_config: WriteMetadataConfig or True (for all defaults)
        source_connection: Name of the source connection
        source_table: Name of the source table (SQL sources)
        source_path: Path of the source file (file sources)
        is_file_source: True if source is a file-based read

    Returns:
        DataFrame with metadata columns added
    """
    from odibi.config import WriteMetadataConfig

    if metadata_config is True:
        config = WriteMetadataConfig()
    elif isinstance(metadata_config, WriteMetadataConfig):
        config = metadata_config
    else:
        return df

    columns = []
    if config.extracted_at:
        columns.append(pl.lit(datetime.now(timezone.utc)).alias("_extracted_at"))
    if config.source_file and is_file_source and source_path:
        columns.append(pl.lit(source_path).alias("_source_file"))
    if config.source_connection and source_connection:
        columns.append(pl.lit(source_connection).alias("_source_connection"))
    if config.source_table and source_table:
        columns.append(pl.lit(source_table).alias("_source_table"))

    if columns:
        df = df.with_columns(columns)

    return df

anonymize(df, columns, method, salt=None)

Anonymize specified columns.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to anonymize.

required
columns List[str]

List of column names to anonymize.

required
method str

Anonymization method ('mask', 'hash', 'redact').

required
salt Optional[str]

Optional salt string for hash-based anonymization.

None

Returns:

Type Description
Any

DataFrame or LazyFrame with anonymized columns.

Source code in odibi/engine/polars_engine.py
def anonymize(
    self, df: Any, columns: List[str], method: str, salt: Optional[str] = None
) -> Any:
    """Anonymize specified columns.

    Args:
        df: DataFrame or LazyFrame to anonymize.
        columns: List of column names to anonymize.
        method: Anonymization method ('mask', 'hash', 'redact').
        salt: Optional salt string for hash-based anonymization.

    Returns:
        DataFrame or LazyFrame with anonymized columns.
    """
    if method == "mask":
        # Mask all but last 4 characters: '******1234'
        # Regex look-around not supported in some envs.
        # Manual approach:
        # If len > 4: repeat('*', len-4) + suffix(4)
        # Else: keep original (or mask all? Pandas engine masked all but last 4, which implies keeping small strings?)
        # Pandas: .str.replace(r".(?=.{4})", "*") -> replaces chars that are followed by 4 chars.
        # If str is "123", no char is followed by 4 chars -> "123".
        # If str is "12345", '1' is followed by '2345' (4 chars) -> "*2345".

        return df.with_columns(
            [
                pl.when(pl.col(c).cast(pl.Utf8).str.len_chars() > 4)
                .then(
                    pl.concat_str(
                        [
                            # Cast to signed Int64 BEFORE subtracting, then clip to 0:
                            # len_chars() is UInt32, so `len - 4` for short strings would
                            # underflow to a huge unsigned value. Polars evaluates .then()
                            # for ALL rows (not just len>4), so without this repeat_by would
                            # try to allocate tens of GB and abort the process.
                            pl.lit("*")
                            .repeat_by(
                                (pl.col(c).str.len_chars().cast(pl.Int64) - 4).clip(
                                    lower_bound=0
                                )
                            )
                            .list.join(""),
                            pl.col(c).str.slice(-4),
                        ]
                    )
                )
                .otherwise(pl.col(c).cast(pl.Utf8))
                .alias(c)
                for c in columns
            ]
        )

    elif method == "hash":
        # Polars hash() is non-cryptographic usually (xxHash).
        # For cryptographic hash (sha256), we might need map_elements (slow) or plugin.
        # Requirement is just 'hash', often consistent for analytics.
        # Gap Analysis mentions "salt".
        # PandasEngine used sha256 with salt.
        # Polars `hash` is fast 64-bit hash.
        # If we need SHA256, we must use map_elements (python UDF) or custom.
        # For "High Performance", map_elements is bad.
        # However, without native plugin, we have no choice for SHA256.
        # Let's implement SHA256 via map_elements for compatibility,
        # OR use Polars internal hash if user accepts non-crypto.
        # But "salt" implies security/crypto usage.

        def _hash_val(val):
            if val is None:
                return None
            to_hash = str(val)
            if salt:
                to_hash += salt
            return hashlib.sha256(to_hash.encode("utf-8")).hexdigest()

        # Apply to each column. Warning: Slow path.
        # But Polars UDFs are still faster than Pandas apply often due to no GIL? No, Python UDF has GIL.
        return df.with_columns(
            [pl.col(c).map_elements(_hash_val, return_dtype=pl.Utf8).alias(c) for c in columns]
        )

    elif method == "redact":
        return df.with_columns([pl.lit("[REDACTED]").alias(c) for c in columns])

    return df

count_nulls(df, columns)

Count nulls in specified columns.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to analyze.

required
columns List[str]

List of column names to count nulls for.

required

Returns:

Type Description
Dict[str, int]

Dictionary mapping column names to null counts.

Source code in odibi/engine/polars_engine.py
def count_nulls(self, df: Any, columns: List[str]) -> Dict[str, int]:
    """Count nulls in specified columns.

    Args:
        df: DataFrame or LazyFrame to analyze.
        columns: List of column names to count nulls for.

    Returns:
        Dictionary mapping column names to null counts.
    """
    if isinstance(df, pl.LazyFrame):
        # efficient null count
        return df.select([pl.col(c).null_count() for c in columns]).collect().to_dicts()[0]

    return df.select([pl.col(c).null_count() for c in columns]).to_dicts()[0]

count_rows(df)

Count rows in DataFrame.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to count rows from.

required

Returns:

Type Description
int

Number of rows as integer.

Source code in odibi/engine/polars_engine.py
def count_rows(self, df: Any) -> int:
    """Count rows in DataFrame.

    Args:
        df: DataFrame or LazyFrame to count rows from.

    Returns:
        Number of rows as integer.
    """
    if isinstance(df, pl.LazyFrame):
        return df.select(pl.len()).collect().item()
    return len(df)

execute_operation(operation, params, df)

Execute built-in operation.

Parameters:

Name Type Description Default
operation str

Name of the operation to execute (e.g., 'pivot', 'unpivot', 'explode').

required
params Dict[str, Any]

Dictionary of parameters specific to the operation.

required
df Any

DataFrame or LazyFrame to operate on.

required

Returns:

Type Description
Any

Transformed DataFrame or LazyFrame.

Source code in odibi/engine/polars_engine.py
def execute_operation(self, operation: str, params: Dict[str, Any], df: Any) -> Any:
    """Execute built-in operation.

    Args:
        operation: Name of the operation to execute (e.g., 'pivot', 'unpivot', 'explode').
        params: Dictionary of parameters specific to the operation.
        df: DataFrame or LazyFrame to operate on.

    Returns:
        Transformed DataFrame or LazyFrame.
    """
    # Ensure LazyFrame for consistency if possible, but operations work on both usually.
    # If DataFrame, some operations might need different methods.

    if operation == "pivot":
        # Pivot requires materialization usually in other engines, but Polars LazyFrame has 'collect' or similar constraints?
        # Polars lazy pivot is not fully supported in older versions without collect, but check recent.
        # Pivot changes shape drastically.
        # params: pivot_column, value_column, group_by, agg_func

        # If lazy, we might need to collect for pivot if lazy pivot isn't supported or experimental.
        # But let's try to keep it lazy if possible.
        # As of recent Polars, pivot is available on DataFrame, experimental on LazyFrame?
        # Actually, 'unstack' or 'pivot' on LazyFrame is limited.
        # Safe bet: materialize if needed, or use lazy pivot if available.

        # Let's collect if input is lazy, because pivot usually implies strict schema change hard to predict.
        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        return df.pivot(
            index=params.get("group_by"),
            on=params["pivot_column"],
            values=params["value_column"],
            aggregate_function=params.get("agg_func", "first"),
        )  # Returns DataFrame

    elif operation == "drop_duplicates":
        subset = params.get("subset")
        if isinstance(df, pl.LazyFrame):
            return df.unique(subset=subset)
        return df.unique(subset=subset)

    elif operation == "fillna":
        value = params.get("value")
        # Polars uses fill_null
        if isinstance(value, dict):
            # Fill specific columns
            # value = {'col1': 0, 'col2': 'unknown'}
            # We need to chain with_columns
            exprs = []
            for col, val in value.items():
                exprs.append(pl.col(col).fill_null(val))
            return df.with_columns(exprs)
        else:
            # Fill all columns? Polars fill_null requires specifying columns or using all()
            return df.fill_null(value)

    elif operation == "drop":
        columns = params.get("columns") or params.get("labels")
        return df.drop(columns)

    elif operation == "rename":
        columns = params.get("columns") or params.get("mapper")
        return df.rename(columns)

    elif operation == "sort":
        by = params.get("by")
        descending = not params.get("ascending", True)
        if isinstance(df, pl.LazyFrame):
            return df.sort(by, descending=descending)
        return df.sort(by, descending=descending)

    elif operation == "sample":
        # Sample n or frac
        n = params.get("n")
        frac = params.get("frac")
        seed = params.get("random_state")

        # Lazy sample supported
        if n is not None:
            # Note: Polars Lazy sample might be approximate or require 'collect' depending on version/backend?
            # But usually supported.
            if isinstance(df, pl.LazyFrame):
                # LazyFrame.sample takes n (int) or fraction.
                # But polars 0.19+ changed sample signature?
                # It's generally `sample(n=..., fraction=..., seed=...)`
                return (
                    df.collect().sample(n=n, seed=seed).lazy()
                )  # Collecting for exact sample n on lazy might be needed if not supported?
                # Actually, fetch(n) is head. Sample is random.
                # Let's materialize for safety with sample as it's often for checks.
                pass
            return df.sample(n=n, seed=seed)
        elif frac is not None:
            if isinstance(df, pl.LazyFrame):
                # Lazy sampling by fraction is supported
                pass  # fall through
            return df.sample(fraction=frac, seed=seed)

    elif operation == "filter":
        # Legacy or simple filter
        pass

    else:
        # Fallback: check if operation is a registered transformer
        from odibi.context import EngineContext, PandasContext
        from odibi.registry import FunctionRegistry

        if FunctionRegistry.has_function(operation):
            func = FunctionRegistry.get_function(operation)
            param_model = FunctionRegistry.get_param_model(operation)

            # Create EngineContext from current df (use PandasContext as placeholder)
            engine_ctx = EngineContext(
                context=PandasContext(),
                df=df,
                engine=self,
                engine_type=self.engine_type,
            )

            # Validate and instantiate params
            if param_model:
                validated_params = param_model(**params)
                result_ctx = func(engine_ctx, validated_params)
            else:
                result_ctx = func(engine_ctx, **params)

            return result_ctx.df

    return df

execute_sql(sql, context)

Execute SQL query using Polars SQLContext.

Parameters:

Name Type Description Default
sql str

SQL query string

required
context Context

Execution context with registered DataFrames

required

Returns:

Type Description
Any

pl.LazyFrame

Source code in odibi/engine/polars_engine.py
def execute_sql(self, sql: str, context: Context) -> Any:
    """Execute SQL query using Polars SQLContext.

    Args:
        sql: SQL query string
        context: Execution context with registered DataFrames

    Returns:
        pl.LazyFrame
    """
    ctx = pl.SQLContext()

    # Register datasets from context
    # We iterate over all registered names in the context
    try:
        names = context.list_names()
        for name in names:
            df = context.get(name)
            # Register LazyFrame or DataFrame
            # Polars SQLContext supports registering LazyFrame, DataFrame, and some others
            # We might need to convert if it's not a Polars object, but we assume Polars engine uses Polars objects
            ctx.register(name, df)
    except Exception:
        # If context doesn't support listing or getting, we proceed with empty context
        # (e.g. if context is not fully compatible or empty)
        pass

    return ctx.execute(sql, eager=False)

filter_coalesce(df, col1, col2, op, value)

Filter using COALESCE(col1, col2) op value.

Automatically casts string columns to datetime for proper comparison.

Parameters:

Name Type Description Default
df DataFrame

Polars DataFrame or LazyFrame

required
col1 str

Primary column name

required
col2 str

Fallback column name

required
op str

Comparison operator (>=, >, <=, <, ==, =)

required
value Any

Value to compare against

required

Returns:

Type Description
DataFrame

Filtered DataFrame

Source code in odibi/engine/polars_engine.py
def filter_coalesce(
    self,
    df: "pl.DataFrame",
    col1: str,
    col2: str,
    op: str,
    value: Any,
) -> "pl.DataFrame":
    """Filter using COALESCE(col1, col2) op value.

    Automatically casts string columns to datetime for proper comparison.

    Args:
        df: Polars DataFrame or LazyFrame
        col1: Primary column name
        col2: Fallback column name
        op: Comparison operator (>=, >, <=, <, ==, =)
        value: Value to compare against

    Returns:
        Filtered DataFrame
    """
    is_lazy = isinstance(df, pl.LazyFrame)
    schema = df.collect_schema() if is_lazy else df.schema

    if col1 not in schema.names():
        raise ValueError(f"Column '{col1}' not found")

    def _cast_if_string(col_name: str) -> pl.Expr:
        if schema[col_name] == pl.Utf8:
            return pl.col(col_name).str.to_datetime(strict=False)
        return pl.col(col_name)

    expr1 = _cast_if_string(col1)
    if col2 in schema.names():
        expr2 = _cast_if_string(col2)
        coalesced = pl.coalesce(expr1, expr2)
    else:
        coalesced = expr1

    try:
        # Determine if value needs datetime parsing by checking schema types
        col1_dtype = schema[col1]
        needs_datetime = False
        if col1_dtype == pl.Utf8:
            needs_datetime = True
        elif col1_dtype in (pl.Datetime, pl.Date):
            needs_datetime = True

        if needs_datetime and isinstance(value, str):
            value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))

        ops = {
            ">=": coalesced >= value,
            ">": coalesced > value,
            "<=": coalesced <= value,
            "<": coalesced < value,
            "==": coalesced == value,
            "=": coalesced == value,
        }

        if op not in ops:
            raise ValueError(f"Unsupported operator: {op}")

        return df.filter(ops[op])
    except Exception as e:
        raise ValueError(f"Failed to filter COALESCE({col1}, {col2}) {op} {value}: {e}")

filter_greater_than(df, column, value)

Filter DataFrame where column > value.

Automatically casts string columns to datetime for proper comparison.

Parameters:

Name Type Description Default
df DataFrame

Polars DataFrame or LazyFrame

required
column str

Column name to filter on

required
value Any

Value to compare against

required

Returns:

Type Description
DataFrame

Filtered DataFrame

Source code in odibi/engine/polars_engine.py
def filter_greater_than(self, df: "pl.DataFrame", column: str, value: Any) -> "pl.DataFrame":
    """Filter DataFrame where column > value.

    Automatically casts string columns to datetime for proper comparison.

    Args:
        df: Polars DataFrame or LazyFrame
        column: Column name to filter on
        value: Value to compare against

    Returns:
        Filtered DataFrame
    """
    is_lazy = isinstance(df, pl.LazyFrame)
    schema = df.collect_schema() if is_lazy else df.schema

    if column not in schema.names():
        raise ValueError(f"Column '{column}' not found in DataFrame")

    try:
        col_dtype = schema[column]

        if col_dtype == pl.Utf8:
            df = df.with_columns(pl.col(column).str.to_datetime(time_unit="us", strict=False))
            value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))
        elif col_dtype in (pl.Datetime, pl.Date) and isinstance(value, str):
            value = datetime.fromisoformat(value.replace("Z", "+00:00").replace(" ", "T"))

        return df.filter(pl.col(column) > value)
    except Exception as e:
        raise ValueError(f"Failed to filter {column} > {value}: {e}")

get_delta_history(connection, path, limit=None)

Get Delta table history.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
limit Optional[int]

Maximum number of versions to return

None

Returns:

Type Description
List[Dict[str, Any]]

List of version metadata dictionaries

Source code in odibi/engine/polars_engine.py
def get_delta_history(
    self, connection: Any, path: str, limit: Optional[int] = None
) -> List[Dict[str, Any]]:
    """Get Delta table history.

    Args:
        connection: Connection object
        path: Delta table path
        limit: Maximum number of versions to return

    Returns:
        List of version metadata dictionaries
    """
    from odibi.utils.logging_context import get_logging_context
    import time

    ctx = get_logging_context().with_context(engine="polars")
    start = time.time()

    ctx.debug("Getting Delta table history", path=path, limit=limit)

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[polars]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path) if connection else path

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    history = dt.history(limit=limit)

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta history retrieved",
        path=str(full_path),
        versions_returned=len(history) if history else 0,
        elapsed_ms=round(elapsed, 2),
    )

    return history

get_sample(df, n=10)

Get sample rows as list of dictionaries.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to sample from.

required
n int

Number of rows to sample (default: 10).

10

Returns:

Type Description
List[Dict[str, Any]]

List of dictionaries, each representing a row.

Source code in odibi/engine/polars_engine.py
def get_sample(self, df: Any, n: int = 10) -> List[Dict[str, Any]]:
    """Get sample rows as list of dictionaries.

    Args:
        df: DataFrame or LazyFrame to sample from.
        n: Number of rows to sample (default: 10).

    Returns:
        List of dictionaries, each representing a row.
    """
    if isinstance(df, pl.LazyFrame):
        return df.limit(n).collect().to_dicts()
    return df.head(n).to_dicts()

get_schema(df)

Get DataFrame schema.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to get schema from.

required

Returns:

Type Description
Any

Dictionary mapping column names to data type strings.

Source code in odibi/engine/polars_engine.py
def get_schema(self, df: Any) -> Any:
    """Get DataFrame schema.

    Args:
        df: DataFrame or LazyFrame to get schema from.

    Returns:
        Dictionary mapping column names to data type strings.
    """
    # Polars schema is a dict {name: DataType}
    # We can return a dict of strings for compatibility
    schema = df.collect_schema() if isinstance(df, pl.LazyFrame) else df.schema
    return {name: str(dtype) for name, dtype in schema.items()}

get_shape(df)

Get DataFrame shape.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to get shape from.

required

Returns:

Type Description
tuple

Tuple of (rows, columns) as integers.

Source code in odibi/engine/polars_engine.py
def get_shape(self, df: Any) -> tuple:
    """Get DataFrame shape.

    Args:
        df: DataFrame or LazyFrame to get shape from.

    Returns:
        Tuple of (rows, columns) as integers.
    """
    if isinstance(df, pl.LazyFrame):
        # Expensive to count rows in LazyFrame without scan
        # But usually shape implies (rows, cols)
        # columns is cheap. rows requires partial scan or metadata.
        # Fetching 1 row might give columns.
        # For exact row count, we need collect(count)
        cols = len(df.collect_schema().names())
        rows = df.select(pl.len()).collect().item()
        return (rows, cols)
    return df.shape

get_source_files(df)

Get list of source files that generated this DataFrame.

For Polars, this checks if source file info was stored in the DataFrame's metadata during read.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame

required

Returns:

Type Description
List[str]

List of file paths (or empty list if not applicable/supported)

Source code in odibi/engine/polars_engine.py
def get_source_files(self, df: Any) -> List[str]:
    """Get list of source files that generated this DataFrame.

    For Polars, this checks if source file info was stored
    in the DataFrame's metadata during read.

    Args:
        df: DataFrame or LazyFrame

    Returns:
        List of file paths (or empty list if not applicable/supported)
    """
    if isinstance(df, pl.LazyFrame):
        return []

    if hasattr(df, "attrs"):
        return df.attrs.get("odibi_source_files", [])

    return []

get_table_schema(connection, table=None, path=None, format=None)

Get schema of an existing table/file.

Parameters:

Name Type Description Default
connection Any

Connection object

required
table Optional[str]

Table name

None
path Optional[str]

File path

None
format Optional[str]

Data format (optional, helps with file-based sources)

None

Returns:

Type Description
Optional[Dict[str, str]]

Schema dict or None if table doesn't exist or schema fetch fails.

Source code in odibi/engine/polars_engine.py
def get_table_schema(
    self,
    connection: Any,
    table: Optional[str] = None,
    path: Optional[str] = None,
    format: Optional[str] = None,
) -> Optional[Dict[str, str]]:
    """Get schema of an existing table/file.

    Args:
        connection: Connection object
        table: Table name
        path: File path
        format: Data format (optional, helps with file-based sources)

    Returns:
        Schema dict or None if table doesn't exist or schema fetch fails.
    """
    from odibi.utils.logging_context import get_logging_context

    ctx = get_logging_context().with_context(engine="polars")

    try:
        if table and is_sql_format(format):
            query = connection.build_select_query(table, limit=0)
            df = connection.read_sql(query)
            return {col: str(dtype) for col, dtype in zip(df.columns, df.dtypes)}

        if path:
            full_path = connection.get_path(path) if connection else path
            if not os.path.exists(full_path):
                return None

            if format == "delta":
                try:
                    from deltalake import DeltaTable

                    dt = DeltaTable(full_path)
                    arrow_schema = dt.schema().to_pyarrow()
                    return {field.name: str(field.type) for field in arrow_schema}
                except ImportError:
                    ctx.warning(
                        "deltalake library not installed for schema introspection",
                        path=full_path,
                    )
                    return None

            elif format == "parquet":
                try:
                    import pyarrow.parquet as pq
                    import glob as glob_mod

                    target_path = full_path
                    if os.path.isdir(full_path):
                        files = glob_mod.glob(os.path.join(full_path, "*.parquet"))
                        if not files:
                            return None
                        target_path = files[0]

                    schema = pq.read_schema(target_path)
                    return {field.name: str(field.type) for field in schema}
                except ImportError:
                    lf = pl.scan_parquet(full_path)
                    schema = lf.collect_schema()
                    return {name: str(dtype) for name, dtype in schema.items()}

            elif format == "csv":
                lf = pl.scan_csv(full_path)
                schema = lf.collect_schema()
                return {name: str(dtype) for name, dtype in schema.items()}

    except (FileNotFoundError, PermissionError):
        return None
    except Exception as e:
        ctx.warning(f"Failed to infer schema for {table or path}: {e}")
        return None

    return None

harmonize_schema(df, target_schema, policy)

Harmonize DataFrame schema.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to harmonize.

required
target_schema Dict[str, str]

Dictionary mapping column names to target data types.

required
policy Any

SchemaPolicyConfig object defining harmonization behavior (mode, on_missing_columns, on_new_columns).

required

Returns:

Type Description
Any

DataFrame or LazyFrame with harmonized schema.

Source code in odibi/engine/polars_engine.py
def harmonize_schema(self, df: Any, target_schema: Dict[str, str], policy: Any) -> Any:
    """Harmonize DataFrame schema.

    Args:
        df: DataFrame or LazyFrame to harmonize.
        target_schema: Dictionary mapping column names to target data types.
        policy: SchemaPolicyConfig object defining harmonization behavior
            (mode, on_missing_columns, on_new_columns).

    Returns:
        DataFrame or LazyFrame with harmonized schema.
    """
    # policy: SchemaPolicyConfig
    from odibi.config import OnMissingColumns, OnNewColumns, SchemaMode

    # Helper to get current columns/schema
    if isinstance(df, pl.LazyFrame):
        current_schema = df.collect_schema()
    else:
        current_schema = df.schema

    current_cols = current_schema.names()
    target_cols = list(target_schema.keys())

    missing = set(target_cols) - set(current_cols)
    new_cols = set(current_cols) - set(target_cols)

    # 1. Validation
    if missing and getattr(policy, "on_missing_columns", None) == OnMissingColumns.FAIL:
        raise ValueError(
            f"Schema Policy Violation: DataFrame is missing required columns {missing}. "
            f"Available columns: {current_cols}. Add missing columns or set on_missing_columns policy."
        )

    if new_cols and getattr(policy, "on_new_columns", None) == OnNewColumns.FAIL:
        raise ValueError(
            f"Schema Policy Violation: DataFrame contains unexpected columns {new_cols}. "
            f"Expected columns: {target_cols}. Remove extra columns or set on_new_columns policy."
        )

    # 2. Transformations
    exprs = []

    # Handle Missing (Add nulls)
    # Evolve means we keep new columns, Enforce means we select only target
    mode = getattr(policy, "mode", SchemaMode.ENFORCE)

    if (
        mode == SchemaMode.EVOLVE
        and getattr(policy, "on_new_columns", None) == OnNewColumns.ADD_NULLABLE
    ):
        # Add missing (if missing cols exist, we fill them with nulls)
        # on_missing_columns controls what to do with missing target cols.
        # If mode is EVOLVE, we typically keep everything?
        # But harmonize_schema is about matching a TARGET schema.
        # If target has cols that df doesn't:
        # If on_missing_columns == FILL_NULL -> Add them as null.
        pass

    # We should respect on_missing_columns regardless of mode?
    if missing and getattr(policy, "on_missing_columns", None) == OnMissingColumns.FILL_NULL:
        for col in missing:
            exprs.append(pl.lit(None).alias(col))

    if exprs:
        df = df.with_columns(exprs)

    # Now Select
    if mode == SchemaMode.ENFORCE:
        # Select only target columns.
        # Missing columns were added above if configured.
        # New columns (not in target) are dropped implicitly by selecting target_cols.
        # But wait, we added exprs to df (lazy).

        final_cols = []
        for col in target_cols:
            final_cols.append(pl.col(col))

        df = df.select(final_cols)

    elif mode == SchemaMode.EVOLVE:
        # We keep new columns.
        # If target has columns that were missing in df, we added them above (if FILL_NULL).
        # If df has columns not in target (new_cols), we keep them.
        pass

    return df

maintain_table(connection, format, table=None, path=None, config=None)

Run table maintenance operations (optimize, vacuum) for Delta tables.

Parameters:

Name Type Description Default
connection Any

Connection object

required
format str

Table format

required
table Optional[str]

Table name

None
path Optional[str]

Table path

None
config Optional[Any]

AutoOptimizeConfig object

None

Returns:

Type Description
None

None

Source code in odibi/engine/polars_engine.py
def maintain_table(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    config: Optional[Any] = None,
) -> None:
    """Run table maintenance operations (optimize, vacuum) for Delta tables.

    Args:
        connection: Connection object
        format: Table format
        table: Table name
        path: Table path
        config: AutoOptimizeConfig object

    Returns:
        None
    """
    from odibi.utils.logging_context import get_logging_context

    ctx = get_logging_context().with_context(engine="polars")

    if format != "delta" or not config or not getattr(config, "enabled", False):
        return

    if not path and not table:
        return

    full_path = connection.get_path(path if path else table) if connection else (path or table)

    ctx.info("Starting table maintenance", path=str(full_path))

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.warning(
            "Auto-optimize skipped: 'deltalake' library not installed",
            path=str(full_path),
        )
        return

    try:
        import time

        start = time.time()

        storage_opts = {}
        if hasattr(connection, "pandas_storage_options"):
            storage_opts = connection.pandas_storage_options()

        dt = DeltaTable(full_path, storage_options=storage_opts)

        ctx.info("Running Delta OPTIMIZE (compaction)", path=str(full_path))
        dt.optimize.compact()

        retention = getattr(config, "vacuum_retention_hours", None)
        if retention is not None and retention > 0:
            ctx.info(
                "Running Delta VACUUM",
                path=str(full_path),
                retention_hours=retention,
            )
            dt.vacuum(
                retention_hours=retention,
                enforce_retention_duration=True,
                dry_run=False,
            )

        elapsed = (time.time() - start) * 1000
        ctx.info(
            "Table maintenance completed",
            path=str(full_path),
            elapsed_ms=round(elapsed, 2),
        )

    except Exception as e:
        ctx.warning(
            "Auto-optimize failed",
            path=str(full_path),
            error=str(e),
        )

materialize(df)

Materialize lazy dataset into memory (DataFrame).

Parameters:

Name Type Description Default
df Any

LazyFrame or DataFrame

required

Returns:

Type Description
Any

Materialized DataFrame (pl.DataFrame)

Source code in odibi/engine/polars_engine.py
def materialize(self, df: Any) -> Any:
    """Materialize lazy dataset into memory (DataFrame).

    Args:
        df: LazyFrame or DataFrame

    Returns:
        Materialized DataFrame (pl.DataFrame)
    """
    if isinstance(df, pl.LazyFrame):
        return df.collect()
    return df

profile_nulls(df)

Calculate null percentage for each column.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to profile.

required

Returns:

Type Description
Dict[str, float]

Dictionary mapping column names to null percentages (0.0 to 1.0).

Source code in odibi/engine/polars_engine.py
def profile_nulls(self, df: Any) -> Dict[str, float]:
    """Calculate null percentage for each column.

    Args:
        df: DataFrame or LazyFrame to profile.

    Returns:
        Dictionary mapping column names to null percentages (0.0 to 1.0).
    """
    if isinstance(df, pl.LazyFrame):
        # null_count() / count()
        # We can do this in one expression
        total_count = df.select(pl.len()).collect().item()
        if total_count == 0:
            return {col: 0.0 for col in df.collect_schema().names()}

        cols = df.collect_schema().names()
        null_counts = df.select([pl.col(c).null_count().alias(c) for c in cols]).collect()
        return {col: null_counts[col][0] / total_count for col in cols}

    total_count = len(df)
    if total_count == 0:
        return {col: 0.0 for col in df.columns}

    null_counts = df.null_count()
    return {col: null_counts[col][0] / total_count for col in df.columns}

read(connection, format, table=None, path=None, streaming=False, schema=None, options=None, **kwargs)

Read data using Polars (Lazy by default).

Parameters:

Name Type Description Default
connection Any

Connection object providing base path and storage options.

required
format str

Data format (csv, parquet, delta, json, sql, sql_server, azure_sql).

required
table Optional[str]

Table name (mutually exclusive with path).

None
path Optional[str]

File path (mutually exclusive with table).

None
streaming bool

Whether to enable streaming mode (uses scan methods when possible).

False
schema Optional[str]

Optional schema specification for SQL queries.

None
options Optional[Dict[str, Any]]

Additional read options to pass to Polars readers.

None
**kwargs

Additional keyword arguments.

{}

Returns:

Type Description
Any

pl.LazyFrame or pl.DataFrame

Source code in odibi/engine/polars_engine.py
def read(
    self,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    streaming: bool = False,
    schema: Optional[str] = None,
    options: Optional[Dict[str, Any]] = None,
    **kwargs,
) -> Any:
    """Read data using Polars (Lazy by default).

    Args:
        connection: Connection object providing base path and storage options.
        format: Data format (csv, parquet, delta, json, sql, sql_server, azure_sql).
        table: Table name (mutually exclusive with path).
        path: File path (mutually exclusive with table).
        streaming: Whether to enable streaming mode (uses scan methods when possible).
        schema: Optional schema specification for SQL queries.
        options: Additional read options to pass to Polars readers.
        **kwargs: Additional keyword arguments.

    Returns:
        pl.LazyFrame or pl.DataFrame
    """
    options = options or {}

    # SQL Server / Azure SQL Support
    if is_sql_format(format):
        if not hasattr(connection, "read_table") and not hasattr(connection, "read_sql_query"):
            raise ValueError(
                f"Cannot read SQL table: connection type '{type(connection).__name__}' "
                "does not support SQL operations. Use a SQL-compatible connection."
            )
        table_name = table or path
        if not table_name:
            raise ValueError("SQL read requires 'table' or 'path' to specify the table name.")
        default_schema = getattr(connection, "default_schema", "dbo")
        if "." in table_name:
            schema_name, tbl = table_name.split(".", 1)
        else:
            schema_name, tbl = default_schema, table_name

        sql_filter = options.get("filter")

        if sql_filter and hasattr(connection, "read_sql_query"):
            full_query = connection.build_select_query(tbl, schema_name, where=sql_filter)
            pdf = connection.read_sql_query(full_query)
        else:
            pdf = connection.read_table(table_name=tbl, schema=schema_name)

        return pl.from_pandas(pdf).lazy()

    # Simulation format: delegate to Pandas, convert to Polars
    if format == "simulation":
        from odibi.engine.pandas_engine import PandasEngine
        from odibi.utils.logging_context import get_logging_context

        ctx = get_logging_context()
        ctx.debug("Reading simulation via Pandas engine, converting to Polars")

        # Use Pandas engine for simulation
        pandas_engine = PandasEngine()
        pdf = pandas_engine._read_simulation(options, ctx)

        # Convert to Polars LazyFrame
        lf = pl.from_pandas(pdf).lazy()

        # Preserve HWM metadata
        if hasattr(pdf, "attrs") and "_simulation_max_timestamp" in pdf.attrs:
            # Polars LazyFrame doesn't have attrs, so we attach as property
            lf._simulation_max_timestamp = pdf.attrs["_simulation_max_timestamp"]

        # Preserve random walk state
        if hasattr(pdf, "attrs") and "_simulation_random_walk_state" in pdf.attrs:
            lf._simulation_random_walk_state = pdf.attrs["_simulation_random_walk_state"]

        # Preserve scheduled event state
        if hasattr(pdf, "attrs") and "_simulation_scheduled_event_state" in pdf.attrs:
            lf._simulation_scheduled_event_state = pdf.attrs[
                "_simulation_scheduled_event_state"
            ]

        ctx.info(
            "Simulation read completed (via Pandas, converted to Polars LazyFrame)",
            row_count=len(pdf),
        )
        return lf

    # Get full path
    if path:
        if connection:
            full_path = connection.get_path(path)
        else:
            full_path = path
    elif table:
        if connection:
            full_path = connection.get_path(table)
        else:
            raise ValueError(
                f"Cannot read table '{table}': connection is required when using 'table' parameter. "
                "Provide a valid connection object or use 'path' for file-based reads."
            )
    else:
        raise ValueError(
            "Read operation failed: neither 'path' nor 'table' was provided. "
            "Specify a file path or table name in your configuration."
        )

    # Handle glob patterns/lists
    # Polars scan methods often support glob strings directly.

    try:
        if format == "csv":
            # scan_csv supports glob patterns
            return pl.scan_csv(full_path, **options)

        elif format == "parquet":
            return pl.scan_parquet(full_path, **options)

        elif format == "json":
            # scan_ndjson for newline delimited json, read_json for standard
            # Assuming ndjson/jsonl for big data usually
            if options.get("json_lines", True):  # Default to ndjson scan
                return pl.scan_ndjson(full_path, **options)
            else:
                # Standard JSON doesn't support lazy scan well in all versions, fallback to read
                return pl.read_json(full_path, **options).lazy()

        elif format == "delta":
            # scan_delta requires 'deltalake' extra usually or feature
            storage_options = options.get("storage_options", None)
            version = options.get("versionAsOf", None)

            # scan_delta is available in recent polars
            # It might accept storage_options in recent versions
            delta_opts = {}
            if storage_options:
                delta_opts["storage_options"] = storage_options
            if version is not None:
                delta_opts["version"] = version

            return pl.scan_delta(full_path, **delta_opts)

        elif format == "excel":
            # Excel format: delegate to Pandas (best Excel support), convert to Polars
            from odibi.engine.pandas_engine import PandasEngine
            from odibi.context import get_logging_context

            ctx = get_logging_context().with_context(engine="polars")
            ctx.debug("Reading Excel via Pandas engine (best Excel support)")

            # Get storage_options from connection for Azure/cloud authentication
            excel_storage_options = None
            if hasattr(connection, "pandas_storage_options"):
                excel_storage_options = connection.pandas_storage_options()

            # Use Pandas engine for Excel reading
            pandas_engine = PandasEngine()
            pdf = pandas_engine._read_excel_with_patterns(
                full_path,
                sheet_pattern=options.pop("sheet_pattern", None),
                sheet_pattern_case_sensitive=options.pop("sheet_pattern_case_sensitive", False),
                add_source_file=options.pop("add_source_file", False),
                is_glob="*" in str(full_path) or "?" in str(full_path),
                ctx=ctx,
                storage_options=excel_storage_options,
                **options,
            )

            # Convert Pandas DataFrame to Polars LazyFrame
            ctx.info(f"Excel read completed (via Pandas): {path}", row_count=len(pdf))
            return pl.from_pandas(pdf).lazy()

        elif format == "api":
            # API format: delegate to Pandas (uses ApiFetcher), convert to Polars
            from odibi.connections.api_fetcher import create_api_fetcher
            from odibi.connections.http import HttpConnection
            from odibi.context import get_logging_context

            ctx = get_logging_context().with_context(engine="polars")
            ctx.debug("Reading API via ApiFetcher", endpoint=str(path))

            if not isinstance(connection, HttpConnection):
                raise ValueError(
                    f"Cannot read API data: connection type '{type(connection).__name__}' "
                    "is not an HttpConnection. Use an HTTP connection for API format."
                )

            # Extract API-specific options
            api_options = options.copy()
            params = api_options.pop("params", {})
            max_records = api_options.pop("max_records", None)
            method = api_options.pop("method", "GET")
            request_body = api_options.pop("request_body", None)

            # Create fetcher with connection's base_url and headers
            fetcher = create_api_fetcher(
                base_url=connection.base_url,
                headers=connection.headers,
                options=api_options,
            )

            # Fetch data as Pandas DataFrame
            pdf = fetcher.fetch_dataframe(
                endpoint=str(full_path),
                params=params,
                max_records=max_records,
                method=method,
                request_body=request_body,
            )

            ctx.info(f"API read completed: {path}", row_count=len(pdf))
            return pl.from_pandas(pdf).lazy()

        else:
            raise ValueError(
                f"Unsupported format for Polars engine: '{format}'. "
                "Supported formats: csv, parquet, json, delta, excel, api, sql, sql_server, azure_sql."
            )

    except Exception as e:
        raise ValueError(
            f"Failed to read {format} from '{full_path}': {e}. "
            "Check that the file exists, the format is correct, and you have read permissions."
        )

restore_delta(connection, path, version)

Restore Delta table to a specific version.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
version int

Version number to restore to

required
Source code in odibi/engine/polars_engine.py
def restore_delta(self, connection: Any, path: str, version: int) -> None:
    """Restore Delta table to a specific version.

    Args:
        connection: Connection object
        path: Delta table path
        version: Version number to restore to
    """
    from odibi.utils.logging_context import get_logging_context

    ctx = get_logging_context().with_context(engine="polars")
    start = time.time()

    ctx.info("Starting Delta table restore", path=path, target_version=version)

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[polars]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path) if connection else path

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    dt.restore(version)

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta table restored",
        path=str(full_path),
        restored_to_version=version,
        elapsed_ms=round(elapsed, 2),
    )

table_exists(connection, table=None, path=None)

Check if table or location exists.

Parameters:

Name Type Description Default
connection Any

Connection object to resolve paths.

required
table Optional[str]

Table name to check (mutually exclusive with path).

None
path Optional[str]

File path to check (mutually exclusive with table).

None

Returns:

Type Description
bool

True if the table or path exists, False otherwise.

Source code in odibi/engine/polars_engine.py
def table_exists(
    self, connection: Any, table: Optional[str] = None, path: Optional[str] = None
) -> bool:
    """Check if table or location exists.

    Args:
        connection: Connection object to resolve paths.
        table: Table name to check (mutually exclusive with path).
        path: File path to check (mutually exclusive with table).

    Returns:
        True if the table or path exists, False otherwise.
    """
    if path:
        full_path = connection.get_path(path)
        return os.path.exists(full_path)
    return False

vacuum_delta(connection, path, retention_hours=168, dry_run=False, enforce_retention_duration=True)

VACUUM a Delta table to remove old files.

Parameters:

Name Type Description Default
connection Any

Connection object

required
path str

Delta table path

required
retention_hours int

Retention period (default 168 = 7 days)

168
dry_run bool

If True, only show files to be deleted

False
enforce_retention_duration bool

If False, allows retention < 168 hours (testing only)

True

Returns:

Type Description
Dict[str, Any]

Dictionary with files_deleted count

Source code in odibi/engine/polars_engine.py
def vacuum_delta(
    self,
    connection: Any,
    path: str,
    retention_hours: int = 168,
    dry_run: bool = False,
    enforce_retention_duration: bool = True,
) -> Dict[str, Any]:
    """VACUUM a Delta table to remove old files.

    Args:
        connection: Connection object
        path: Delta table path
        retention_hours: Retention period (default 168 = 7 days)
        dry_run: If True, only show files to be deleted
        enforce_retention_duration: If False, allows retention < 168 hours (testing only)

    Returns:
        Dictionary with files_deleted count
    """
    from odibi.utils.logging_context import get_logging_context
    import time

    ctx = get_logging_context().with_context(engine="polars")
    start = time.time()

    ctx.debug(
        "Starting Delta VACUUM",
        path=path,
        retention_hours=retention_hours,
        dry_run=dry_run,
    )

    try:
        from deltalake import DeltaTable
    except ImportError:
        ctx.error("Delta Lake library not installed", path=path)
        raise ImportError(
            "Delta Lake support requires 'pip install odibi[polars]' "
            "or 'pip install deltalake'. See README.md for installation instructions."
        )

    full_path = connection.get_path(path) if connection else path

    storage_opts = {}
    if hasattr(connection, "pandas_storage_options"):
        storage_opts = connection.pandas_storage_options()

    dt = DeltaTable(full_path, storage_options=storage_opts)
    deleted_files = dt.vacuum(
        retention_hours=retention_hours,
        dry_run=dry_run,
        enforce_retention_duration=enforce_retention_duration,
    )

    elapsed = (time.time() - start) * 1000
    ctx.info(
        "Delta VACUUM completed",
        path=str(full_path),
        files_deleted=len(deleted_files),
        dry_run=dry_run,
        elapsed_ms=round(elapsed, 2),
    )

    return {"files_deleted": len(deleted_files)}

validate_data(df, validation_config)

Validate data against rules.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame

required
validation_config Any

ValidationConfig object

required

Returns:

Type Description
List[str]

List of validation failure messages

Source code in odibi/engine/polars_engine.py
def validate_data(self, df: Any, validation_config: Any) -> List[str]:
    """Validate data against rules.

    Args:
        df: DataFrame or LazyFrame
        validation_config: ValidationConfig object

    Returns:
        List of validation failure messages
    """
    failures = []

    if isinstance(df, pl.LazyFrame):
        schema = df.collect_schema()
        columns = schema.names()
    else:
        columns = df.columns

    if getattr(validation_config, "not_empty", False):
        count = self.count_rows(df)
        if count == 0:
            failures.append("DataFrame is empty")

    if getattr(validation_config, "no_nulls", None):
        cols = validation_config.no_nulls
        null_counts = self.count_nulls(df, cols)
        for col, count in null_counts.items():
            if count > 0:
                failures.append(f"Column '{col}' has {count} null values")

    if getattr(validation_config, "schema_validation", None):
        schema_failures = self.validate_schema(df, validation_config.schema_validation)
        failures.extend(schema_failures)

    if getattr(validation_config, "ranges", None):
        for col, bounds in validation_config.ranges.items():
            if col in columns:
                min_val = bounds.get("min")
                max_val = bounds.get("max")

                if min_val is not None:
                    if isinstance(df, pl.LazyFrame):
                        min_violations = (
                            df.filter(pl.col(col) < min_val).select(pl.len()).collect().item()
                        )
                    else:
                        min_violations = len(df.filter(pl.col(col) < min_val))
                    if min_violations > 0:
                        failures.append(f"Column '{col}' has values < {min_val}")

                if max_val is not None:
                    if isinstance(df, pl.LazyFrame):
                        max_violations = (
                            df.filter(pl.col(col) > max_val).select(pl.len()).collect().item()
                        )
                    else:
                        max_violations = len(df.filter(pl.col(col) > max_val))
                    if max_violations > 0:
                        failures.append(f"Column '{col}' has values > {max_val}")
            else:
                failures.append(f"Column '{col}' not found for range validation")

    if getattr(validation_config, "allowed_values", None):
        for col, allowed in validation_config.allowed_values.items():
            if col in columns:
                if isinstance(df, pl.LazyFrame):
                    invalid_count = (
                        df.filter(~pl.col(col).is_in(allowed)).select(pl.len()).collect().item()
                    )
                else:
                    invalid_count = len(df.filter(~pl.col(col).is_in(allowed)))
                if invalid_count > 0:
                    failures.append(f"Column '{col}' has invalid values")
            else:
                failures.append(f"Column '{col}' not found for allowed values validation")

    return failures

validate_schema(df, schema_rules)

Validate DataFrame schema.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to validate.

required
schema_rules Dict[str, Any]

Dictionary containing schema validation rules (e.g., required_columns, expected_types, disallowed_columns).

required

Returns:

Type Description
List[str]

List of validation failure messages (empty if all validations pass).

Source code in odibi/engine/polars_engine.py
def validate_schema(self, df: Any, schema_rules: Dict[str, Any]) -> List[str]:
    """Validate DataFrame schema.

    Args:
        df: DataFrame or LazyFrame to validate.
        schema_rules: Dictionary containing schema validation rules
            (e.g., required_columns, expected_types, disallowed_columns).

    Returns:
        List of validation failure messages (empty if all validations pass).
    """
    failures = []

    # Schema is dict-like in Polars
    current_schema = df.collect_schema() if isinstance(df, pl.LazyFrame) else df.schema
    current_cols = current_schema.keys()

    if "required_columns" in schema_rules:
        required = schema_rules["required_columns"]
        missing = set(required) - set(current_cols)
        if missing:
            failures.append(f"Missing required columns: {', '.join(missing)}")

    if "types" in schema_rules:
        for col, expected_type in schema_rules["types"].items():
            if col not in current_cols:
                failures.append(f"Column '{col}' not found for type validation")
                continue

            actual_type = str(current_schema[col])
            # Basic type check - simplistic string matching
            if expected_type.lower() not in actual_type.lower():
                failures.append(
                    f"Column '{col}' has type '{actual_type}', expected '{expected_type}'"
                )

    return failures

write(df, connection, format, table=None, path=None, mode='overwrite', options=None, streaming_config=None)

Write data using Polars.

Parameters:

Name Type Description Default
df Any

DataFrame or LazyFrame to write.

required
connection Any

Connection object providing base path and storage options.

required
format str

Output format (csv, parquet, delta, json, sql, sql_server, azure_sql).

required
table Optional[str]

Table name (mutually exclusive with path).

None
path Optional[str]

File path (mutually exclusive with table).

None
mode str

Write mode (overwrite, append, upsert, append_once).

'overwrite'
options Optional[Dict[str, Any]]

Additional write options to pass to Polars writers.

None
streaming_config Optional[Any]

Streaming configuration (not used in Polars engine).

None

Returns:

Type Description
Optional[Dict[str, Any]]

Optional dict with write statistics for Delta writes, None otherwise.

Source code in odibi/engine/polars_engine.py
def write(
    self,
    df: Any,
    connection: Any,
    format: str,
    table: Optional[str] = None,
    path: Optional[str] = None,
    mode: str = "overwrite",
    options: Optional[Dict[str, Any]] = None,
    streaming_config: Optional[Any] = None,
) -> Optional[Dict[str, Any]]:
    """Write data using Polars.

    Args:
        df: DataFrame or LazyFrame to write.
        connection: Connection object providing base path and storage options.
        format: Output format (csv, parquet, delta, json, sql, sql_server, azure_sql).
        table: Table name (mutually exclusive with path).
        path: File path (mutually exclusive with table).
        mode: Write mode (overwrite, append, upsert, append_once).
        options: Additional write options to pass to Polars writers.
        streaming_config: Streaming configuration (not used in Polars engine).

    Returns:
        Optional dict with write statistics for Delta writes, None otherwise.
    """
    options = options or {}

    if is_sql_format(format):
        return self._write_sql(df, connection, table, mode, options)

    if path:
        if connection:
            full_path = connection.get_path(path)
        else:
            full_path = path
    elif table:
        if connection:
            full_path = connection.get_path(table)
        else:
            raise ValueError(
                f"Cannot write to table '{table}': connection is required when using 'table' parameter. "
                "Provide a valid connection object or use 'path' for file-based writes."
            )
    else:
        raise ValueError(
            "Write operation failed: neither 'path' nor 'table' was provided. "
            "Specify a file path or table name in your configuration."
        )

    is_lazy = isinstance(df, pl.LazyFrame)

    # Handle upsert/append_once for non-Delta formats
    if mode in ["upsert", "append_once"] and format != "delta":
        df, mode = self._handle_generic_upsert(df, full_path, format, mode, options)
        is_lazy = isinstance(df, pl.LazyFrame)
        options = {k: v for k, v in options.items() if k != "keys"}

    # Handle plain append for file formats — Polars file writers overwrite by default,
    # so without read-existing + concat an append would silently lose prior data.
    if mode == "append" and format in ("parquet", "csv", "json"):
        df, mode = self._handle_file_append(df, full_path, format)
        is_lazy = isinstance(df, pl.LazyFrame)

    parent_dir = os.path.dirname(full_path)
    if parent_dir:
        os.makedirs(parent_dir, exist_ok=True)

    if format == "parquet":
        if is_lazy:
            df.sink_parquet(full_path, **options)
        else:
            df.write_parquet(full_path, **options)

    elif format == "csv":
        if is_lazy:
            df.sink_csv(full_path, **options)
        else:
            df.write_csv(full_path, **options)

    elif format == "json":
        if is_lazy:
            df.sink_ndjson(full_path, **options)
        else:
            df.write_ndjson(full_path, **options)

    elif format == "delta":
        if is_lazy:
            df = df.collect()

        storage_options = options.get("storage_options", None)
        delta_write_options = options.copy()
        if "storage_options" in delta_write_options:
            del delta_write_options["storage_options"]

        df.write_delta(
            full_path, mode=mode, storage_options=storage_options, **delta_write_options
        )

    else:
        raise ValueError(
            f"Unsupported write format for Polars engine: '{format}'. "
            "Supported formats: csv, parquet, json, delta."
        )

    return None